从python中的N-K矩阵循环穿过K-uplets

Question

我必须在python脚本中使用动态编程。

我用shape =（N，K）定义了一个numpy数组u。 我想为每列选择一个元素，因此生成K-uplets。

如何有效地循环以这种方式生成的所有K-uplet？解决方案是使用

import itertools
itertools.combination_with_replacement(list,K) 

其中list = [0..N-1]，但我需要使用itertools方法的输出（索引）迭代地构建每个K-uplet。 还有更直接的方法吗？

由于

文森特

Answer 1

您可以使用arr[ind, np.arange(K)]构建K-uplet。当然，这实际上是一个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]
    # ...