在python中从NK矩阵跨K-uplet循环

Question

我必须在python脚本中使用动态编程。我用shape =（N，K）定义了一个numpy数组u。我想为每一列选择一个元素，因此生成一个K-uplet。您将如何在这种方式生成的所有K-uplet上高效地循环？一个解决方案是使用import itertoolsitertools.combination_with_replacement(list,K) 其中list = [0..N-1]，但我将需要使用itertools方法的输出（索引）来迭代地构建每个K-uplet。是否有更直接的方式进行？

Answer 1

您可以使用构建K-uplet arr[ind, np.arange(K)]。当然，这实际上是一个NumPy ndarray，但是如果您真的想要连音符，则很容易将它们转换为连音符tuple(arr[ind, np.arange(K)])。

import numpy as np

import itertools as IT

N, K = 5,3

arr = np.arange(N*K).reshape(N,K)

print(arr)

# [[ 0  1  2]

#  [ 3  4  5]

#  [ 6  7  8]

#  [ 9 10 11]

#  [12 13 14]]

for ind in IT.combinations_with_replacement(range(N), K):

    print(arr[ind, np.arange(K)])

    # [0 1 2]

    # [0 1 5]

    # [0 1 8]

    # [ 0  1 11]

    # ...