无循环生成满足指定条件的K×2 numpy二元矩阵所有组合方法

实现逻辑 #

我们要生成的所有矩阵，本质是从K行中选2个不同的行，分别在第0列和第1列放1的所有不重复组合，总组合数为$C(K,2)=K*(K-1)/2$，刚好匹配输出数组最后一维的长度要求。全程用NumPy原生操作实现，无显式Python循环，大K值下性能拉满。 #

代码实现 #

import numpy as np

def generate_target_matrices(K):
    # 生成所有i<j的行索引对，对应所有合法组合
    row_i, row_j = np.triu_indices(K, k=1)
    sample_count = len(row_i)
    # 初始化全0数组，形状符合K×2×样本数的要求
    result = np.zeros((K, 2, sample_count), dtype=np.int8)
    # 批量给每个样本的对应位置赋值1
    result[row_i, 0, np.arange(sample_count)] = 1
    result[row_j, 1, np.arange(sample_count)] = 1
    return result

验证示例 #

以K=5为例测试，输出完全匹配你给出的样例：

test_arr = generate_target_matrices(5)
# 输出维度验证：(5, 2, 10)，10=5*4/2符合要求
print(test_arr.shape)

# 查看第1个样本
print(test_arr[..., 0])
# 输出：
# [[1 0]
#  [0 1]
#  [0 0]
#  [0 0]
#  [0 0]]

# 查看第2个样本
print(test_arr[..., 1])
# 输出：
# [[1 0]
#  [0 0]
#  [0 1]
#  [0 0]
#  [0 0]]

性能说明 #

K=300时，总样本量为44850，数组总元素数不到2700万，用int8存储仅占用27MB左右内存，运行时间为毫秒级，完全满足大K值的效率要求。

内容的提问来源于stack exchange，提问作者mariottidae