条件を満たす(n)個の要素を持つリストを全パターン生成したいです。

Question

やりたいこと

• pythonで 条件を満たす(n)個の要素を持つリストを全パターン生成したいです。

例えば、以下のような結果を取得したいです。

 [[1, 1, 1, 1, 1, 1, 1, 2, 2, 2], 
  [0, 1, 1, 1, 1, 1, 1, 2, 2, 2], 
  [0, 0, 1, 1, 1, 1, 1, 2, 2, 2], 
  .....]

生成するリストの条件

• 要素数 n=10 （実際には数を50個以上に増やしたい）
• 各要素は0,1,2からなります。
• 0は3個以下しか含まれない。(nの数が大きくなると、この数字も増やします。例えば、int(n*0.1)個程度。）
• 2も3個以下しか含まれない。（同上）

n=10の場合、3**10パターンから、0と2を一定数以上含むパターンのリストを削除するイメージです。

n=15程度なら、実現できる処理は以下のように書きました。

from itertools import product

pattern = [] 
n = 10 # 何列の結果を出すか（n=20くらいでメモリーエラーになる。）

for _ in range(n):
    pattern.append([0,1,2])

data = list(product(*pattern)) # この処理が重いです。

# dataは、例えば、例： [[0, 0, 1, 1, 1, 1, 1, 2, 2, 2], [0, 1, 1, 1, 1, 1, 1, 2, 2, 2].....]のようなデータです。
# dataから（ここでは例として）、"0が3個以上"、"2が3個以上"含むリストを、削除します。

result = []

for row in data:
    if row.count(0) >= 3 or row.count(2) >= 3:
        continue
    result.append(row)

print(resutl)

この方法ですと、nが大きくなると計算しきれずメモリーエラーになってしまいます。n=10(310)ではすぐに計算できても、20では(320)で手元のPCではエラーを起こします。

ただ、例えば、n=10で計算した場合に、総当りで出したdataは59049個のなのに対して、必要なデータをピックアップしたresultは2181個と相当少なくなります。なので、総当たりせずに、resultを直接出せれば、nが大きくなってもある程度なら計算できるのでは？と考えています。

この処理の中で最も重いのは、

data = list(product(*pattern))

です。
ここで総当りしてしまっているのが無駄な作業で、0や2を含む数が規定より多い場合に、処理をしないようにスキップできればいいはずです。

仮説

例えば、

• はじめに[1,1,1,1,1,1...]を作って、0をn個、2をm個、代入する、という発想もありなのかな？と考えはしたものの、コードに落とし込めません。。
• ビット演算を使う？？
• numpyで処理すれば多少は早くなる？

ただ、正解にたどり着けません。。。

教えて頂けると幸いです。

Answer 1

勝たんしか深さ優先探索な問題ですね．len(cond)分木を探索するイメージで，次のように再帰を書けば余計な配列を生成せず全列挙できます．

n = 10

cond = {
    0: 3,   # 3個未満
    1: 1e9, # 10^9未満(実質無制限)
    2: 3,   # 3個未満
}

result = list()

def dfs(x):
    if len(x) == n:
        result.append(x.copy())
        return
    
    for k, v in cond.items():
        if x.count(k) < v - 1:
            x.append(k)
            dfs(x)
            x.pop()

dfs(list())
print(result)
print(len(result))

n = 10のとき，同様に2181個観測され，n = 30で190591個，n = 50は1504401個の配列が生成されました．適宜，適用したい条件condを追加して実行することができます．

ちなみに生成する条件で「0,2は3個以下」と言っていましたが，3個未満でないと結果の個数が一致しません．どちらが望んだ結果なのか不明でしたので，とりあえず得られる配列の個数に合わせましたが，もし3個以下の条件が正しいなら，if x.count(k) < v - 1:をif x.count(k) < v:に変更してください．

can110さんの回答との相違点は，「グローバル変数を使いまくっているので行儀が悪い」「ファイル書き込みではなくresultに答えを格納しまくるのでメモリ圧迫する」という点です．ファイル書き込みによるメモリ節約に関しては弊回答にも適用しやすいものかと思われます．

Comments

[keiji_kc]

素晴らしい、、、感動しました。

これが深さ優先探索というものでしょうか？大変興味を持ちました。。

お恥ずかしながら、初見で理解できないので、じっくり勉強したいと思います！

手始めにこのあたりを勉強してみようと思います。。
https://qiita.com/drken/items/4a7869c5e304883f539b

[ps_aux_grep]

そうですね，けんちょんさんの資料は大変読みやすいので良い選択だと思います．
深さ優先探索の中でも簡単なものなので，理解できれば非常に大きな一歩となるでしょう．

[keiji_kc]

追記も頂いてありがとうございます。たしかに、一定の数を超えるとメモリエラー出ました。
ファイルに書き出せば解消できるかもしれませんね！ありがとうございます。

Answer 2

max0, max2 = 0および2の最大個数
cnt_total = 合計個数
とすると
0の個数は0～max0
2の個数は0～max2
の範囲をとり、
1の数はcnt_total - cnt0 - cnt2
なので 0および2の個数をそれぞれ0～maxまで変えて
それぞれの個数での組み合わせを列挙していけばよいです。

以下ではとりあえず個数の組み合わせの列挙を再帰で行っていますが、30以上は厳しいと思います。
もっと効率の良い方法あるいはライブラリがあるかもしれません。

import itertools

# 各値の組み合わせを列挙
def count( f, counts, ret=[]):

    # すべて使い切った
    if not any(counts):
        f.write(f'{ret}\n')
        return

    for i in range(len(counts)):
        if counts[i] > 0:
            counts_next = counts.copy()
            counts_next[i] -= 1 # この数(i)を使った
            count( f, counts_next, ret + [i])

cnt_total = 4
max0, max2 = 1,1

with open('ret.txt', 'w') as f:
    for cnt0, cnt2 in itertools.product(range(max0+1), range(max2+1)):
        cnt1 = cnt_total - cnt0 - cnt2
        print(cnt0,cnt1,cnt2) # 0,1,2の個数
        count( f, [cnt0,cnt1,cnt2])

"""
cnt_total = 4
max0, max2 = 1,1
-----
[1, 1, 1, 1]
[1, 1, 1, 2]
[1, 1, 2, 1]
[1, 2, 1, 1]
[2, 1, 1, 1]
[0, 1, 1, 1]
[1, 0, 1, 1]
[1, 1, 0, 1]
[1, 1, 1, 0]
[0, 1, 1, 2]
[0, 1, 2, 1]
[0, 2, 1, 1]
[1, 0, 1, 2]
[1, 0, 2, 1]
[1, 1, 0, 2]
[1, 1, 2, 0]
[1, 2, 0, 1]
[1, 2, 1, 0]
[2, 0, 1, 1]
[2, 1, 0, 1]
[2, 1, 1, 0]
-----
21行
"""

"""
cnt_total = 30
max0, max2 = 3,3
-----
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2]
：
[2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0]
[2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
-----
15143571行
"""

追記

C++で総数だけ数える雑なコード書いてみましたが、40程度でも数分かかって6901580871通りとはじき出されたので、50個ものパターンをじっさいに生成するのは現実的ではないと思います。

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <vector>

using namespace std;

const int TOTAL = 40;
const int MAX0 = 4;
const int MAX2 = 4;

long long nCount = 0;

void count( int aCount[], int aRet[])
{
    if( aCount[0] <= 0 && aCount[1] <= 0 && aCount[2] <= 0){
        nCount++;
        return;
    }

    int pos = TOTAL - (aCount[0] + aCount[1] + aCount[2]);
    for( int i = 0; i < 3; i++){
        if( aCount[i] > 0){
            int aCount_next[3];
            memcpy( aCount_next, aCount, sizeof(aCount_next));
            aCount_next[i]--;
            aRet[pos] = i;
            count(aCount_next, aRet);
        }
    }
}

int main( int argc, char *argv[])
{
    int aRet[TOTAL];
    for( int n0 = 0; n0 <= MAX0; n0++){
        for( int n2 = 0; n2 <= MAX2; n2++){
            int n1 = TOTAL - n0 - n2;
            int aCount[3] = {n0, n1, n2};
            printf("%d,%d,%d\n", aCount[0], aCount[1], aCount[2]);
            count( aCount, aRet);
            printf("%lld\n", nCount);
        }
    }

    return 0;
}

Comments

[keiji_kc]

ありがとうございます。
ファイルに書き出すアイデア、なるほどと思いました！
大変参考になります！

[8524ba23]

ためしにC++で総数を数えるコード書いてみましたが、パターン数が膨大になるのでファイルに書き出すのも厳しいかと思います。

[keiji_kc]

ご丁寧にありがとうございます！
私も教えて頂いて色々やってますが、一定の数を超えると、どうしても限界が来ることはわかりました。

とはいえ、元のコードよりかなり改善できそうでしたので、テストを分割するなど、少し工夫して実現しようと思います。本当にありがとうございました。