如何将二维NumPy数组的所有行对（允许重复、无序）转为三维数组

问题描述 #

给定二维NumPy数组：

import numpy as np
a = np.array([[8, 16, 4, 1, 0, 5],
              [3, 0, 0, 11, 9, 7],
              [5, 5, 8, 5, 15, 5],
              [2, 0, 2, 14, 2, 0],
              [0, 1, 2, 3, 4, 15]])

需要提取所有允许重复行、顺序无关的行对，组合成三维数组。示例中5行的情况下，共有5*(5+1)/2 = 15种组合，最终得到形状为(15, 2, 6)的三维数组，形式如下（示例片段）：

[[[8, 16, 4, 1, 0, 5],
  [8, 16, 4, 1, 0, 5]],
 [[8, 16, 4, 1, 0, 5],
  [3, 0, 0, 11, 9, 7]],
 ...
 [[2, 0, 2, 14, 2, 0],
  [0, 1, 2, 3, 4, 15]],
 [[0, 1, 2, 3, 4, 15],
  [0, 1, 2, 3, 4, 15]]]

要求实现方案高效，适用于大规模数组。

高效实现方案 #

方法：利用np.triu_indices生成索引（推荐） #

np.triu_indices可以直接生成上三角（含对角线）的索引对，正好对应无序且允许重复的行组合，完全避免循环，效率极高。

代码实现：

import numpy as np

a = np.array([[8, 16, 4, 1, 0, 5],
              [3, 0, 0, 11, 9, 7],
              [5, 5, 8, 5, 15, 5],
              [2, 0, 2, 14, 2, 0],
              [0, 1, 2, 3, 4, 15]])

# 获取上三角（含对角线）的索引对
i, j = np.triu_indices(a.shape[0])
# 按索引提取行对，组合成三维数组
result = np.stack([a[i], a[j]], axis=1)

print(result.shape)  # 输出 (15, 2, 6)，符合预期

原理说明 #

• np.triu_indices(n)返回两个长度为n*(n+1)/2的数组，分别代表行索引的上三角位置，正好覆盖所有(i,j)且i<=j的组合，对应无序行对（因为(i,j)和(j,i)视为同一组，这里只取i<=j的情况）。
• 通过a[i]和a[j]直接批量提取行，再用np.stack在axis=1维度合并，一步得到目标三维数组，时间复杂度为O(N²)（N为行数），是NumPy原生的矢量操作，比Python循环快几个数量级，适合大规模数组。

验证结果 #

输出的result就是目标格式，比如前几个元素：

print(result[:3])
# 输出：
[[[ 8 16  4  1  0  5]
  [ 8 16  4  1  0  5]]

 [[ 8 16  4  1  0  5]
  [ 3  0  0 11  9  7]]

 [[ 8 16  4  1  0  5]
  [ 5  5  8  5 15  5]]]

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