python - 从python中的NK矩阵遍历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-uplets。有没有更直接的方法？

谢谢

文森特

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]
    # ...