python - 用 0 和 1 填充的 numpy 二维数组的所有组合

Question

给定 K，我需要拥有 K x 2 numpy 矩阵的所有可能组合，以便在每个矩阵中除了不同行和列中的两个 1 外，其他都是 0。 对于 K = 5 是这样的:

1. [[1,0],[0,1],[0,0],[0,0][0,0]]
2. [[1,0],[0,0],[0,1],[0,0][0,0]]
3. [[1,0],[0,0],[0,0],[0,1][0,0]]
4. [[1,0],[0,0],[0,0],[0,0][0,1]]
5. [[0,0],[1,0],[0,1],[0,0][0,0]]
6. [[0,0],[1,0],[0,0],[0,1][0,0]]
7. ...等等

因此生成的数组应该是 K x 2 x (K*(K-1)/2)。 我想避免循环，因为当 K 足够大时(在我的特定情况下 K = 300)，这不是一种有效的方法

Answer 1

我想不出一个优雅的解决方案，但这里有一个不太优雅的纯 numpy 解决方案:

import numpy as np

def combination_matrices(K):
    # get combination indices
    i, j = np.indices((K, K))
    comb_indices = np.transpose((i < j).nonzero())  # (num_combs, 2) array where ones are
    num_combs = comb_indices.shape[0]  # K*(K-1)/2

    # create a matrix of the desired shape, first axis enumerates combinations
    matrices = np.zeros((num_combs, K, 2), dtype=int)
    # broadcasting assignment of ones
    comb_range, col_index = np.ogrid[:num_combs, :2]
    matrices[comb_range, comb_indices, col_index] = 1
    return matrices

这首先使用 (K, K) 形数组的索引来查找每个组合的索引对(这些索引编码数组的上三角，不包括对角线) .然后我们使用有点棘手的广播赋值(重 fancy indexing )将预分配输出数组的每个对应元素设置为 1。

请注意，我将 K*(K-1)/2 大小的轴放在第一位，因为这在具有 C 连续内存布局的 numpy 中最有意义。这样，当您采用组合索引矩阵 3 时，arr[3, ...] 将是形状为 (K, 2 ) 在矢量化操作中可以快速处理。

K = 4 的输出:

[[[1 0]
  [0 1]
  [0 0]
  [0 0]]

 [[1 0]
  [0 0]
  [0 1]
  [0 0]]

 [[1 0]
  [0 0]
  [0 0]
  [0 1]]

 [[0 0]
  [1 0]
  [0 1]
  [0 0]]

 [[0 0]
  [1 0]
  [0 0]
  [0 1]]

 [[0 0]
  [0 0]
  [1 0]
  [0 1]]]