特殊なリストを高速で作る方法(Pythonで)
Q&A
Closed
以下のような問題を考えます
長さNのリストの中に1がk個, 0がn-k個あるようなものをすべて出力せよ
例:n=2,k=1ならば [0,1],[1,0]
例:n=4,k=3ならば[1,1,1,0],[1,1,0,1],[1,0,1,1],[0,1,1,1]
個数を求めるなら高校でおなじみのコンビネーションにぶち込めば終わりですが, これを高速でやる方法は何がありますか?
0
4Answer
なんの工夫もない総当りプログラムを置いておきます。
n = 10
k = 5
print([list(map(int, list(f'{i:b}'.zfill(n)))) for i in range(1 << n) if bin(i).count('1') == k])
2Like
from more_itertools import distinct_permutations
n, k = 4, 3
list(distinct_permutations([1] * k + [0] * (n - k)))
2Like
pythonは詳しくないので、アルゴリズムだけ書かせていただくことをお許しください。
課題は「nビットの整数で、1が立っているビットがk個のものを全て列挙せよ」と言い換えることができます。
あるいは、「range(0,n)から重複なしにk個取る組み合わせ」と言えます。
(その際、取り出したk個の数値は1が立っているビットを表します。)
組み合わせの算出法については、以下の記事をどうぞ。
Pythonで順列と組み合わせ (Qiita)
Pythonで順列・組み合わせを求める (Qiita)
蛇足 (2020/11/07追記)
勝手に解説
# Carteletさんのコード
n = 10
k = 5
print([list(map(int, list(f'{i:b}'.zfill(n)))) for i in range(1 << n) if bin(i).count('1') == k])
# nビットの全整数中、1の立ったビットがk個の二進数のビットごとのリストのリスト
# 内包表記: リストの内容を式で表現したもの
print([i + 1 for i in range(30) if i % 3 == 0]) # 30未満の3で割りきれる整数に1を加えたリスト
# => [1, 4, 7, 10, 13, 16, 19, 22, 25, 28]
# 上と等価
x = [] # xを空のリストで初期化
for i in range(30): # 0～29まで繰り返す
if i % 3 == 0: # 3で割り切れたら
x.append(i + 1) # xの最後にi+1を加える
print(x)
# => [1, 4, 7, 10, 13, 16, 19, 22, 25, 28]
# 二進文字列化
print(f'{10:b}')
# => 1010
# 二進数化
print(bin(10))
# => 0b1010
# ゼロ埋め
print('x'.zfill(5)) # 'x'に対して、5文字幅で左に'0'でパディングする
# => 0000x
# 組み込み関数map
print(list(map(lambda i : int(i), list(range(5))))) # 整数にする処理を5未満の整数列に対して繰り返したリスト
# => [0, 1, 2, 3, 4]
print(list(map(int, list(range(5))))) # 上と等価
# => [0, 1, 2, 3, 4]
# 出現回数を数える
print(bin(10).count('0')) # 10の二進数(0b1010)から'0'を数える
# => 3
1Like
NumPyを使った方法です。
np.indices()
import numpy as np
def make_array(n, k):
out = np.indices([2]*n).reshape(n, -1).T
return out[out.sum(1) == k]
動作の例
In [3]: make_array(2, 1)
Out[3]:
array([[0, 1],
[1, 0]])
In [4]: make_array(4, 2)
Out[4]:
array([[0, 0, 1, 1],
[0, 1, 0, 1],
[0, 1, 1, 0],
[1, 0, 0, 1],
[1, 0, 1, 0],
[1, 1, 0, 0]])
In [5]: make_array(4, 3)
Out[5]:
array([[0, 1, 1, 1],
[1, 0, 1, 1],
[1, 1, 0, 1],
[1, 1, 1, 0]])
私の環境で計測したところ、$n \geqq 6$の場合は上記の方法のほうが@Cartelet 氏の方法よりも速くなりました。
more_itertools.distinct_permutations()kの値によっては大きく速度が変わるようです。
import benchit
import numpy as np
from more_itertools import distinct_permutations
def BruteForce(n, k):
# @Cartelet
return [list(map(int, list(f'{i:b}'.zfill(n))))
for i in range(1 << n) if bin(i).count('1') == k]
def NumPy(n, k):
out = np.indices([2]*n).reshape(n, -1).T
return out[out.sum(1) == k]
def Itertools(n, k):
# @SaitoTsutomu
return list(distinct_permutations([1] * k + [0] * (n - k)))
funcs = [BruteForce, NumPy, Itertools]
inputs = {(n, k_percent): (n, int(n * k_percent // 100))
for n in range(1, 11) for k_percent in [0, 12.5, 25, 50]}
t = benchit.timings(funcs, inputs, multivar=True, input_name=['n', 'k_percent'])
t.plot(figsize=(6, 4), sp_ncols=2, sp_sharey='g')
1Like