如何对NumPy二维数组进行整行全排列并生成三维数组

直接用itertools.permutations生成行索引的全排列，再结合NumPy的高级索引就能轻松实现需求，代码示例如下：

import numpy as np
from itertools import permutations

# 定义原始数组
arr = np.array([
    [1, 5, 6, 3, 3, 7],
    [2, 2, 2, 2, 2, 2],
    [0, 1, 0, 1, 0, 1],
    [4, 8, 4, 8, 4, 8],
    [1, 2, 3, 4, 5, 6]
])

# 生成所有行索引的全排列（5行共5! = 120种）
row_perms = permutations(range(arr.shape[0]))
# 将排列转换为NumPy数组，形状为(120, 5)
perm_indices = np.array(list(row_perms))
# 通过索引获取所有整行排列的结果，最终形状为(120, 5, 6)
p = arr[perm_indices]

# 验证形状
print(p.shape)  # 输出 (120, 5, 6)

关键步骤说明 #

• 生成行索引排列：permutations(range(5))会生成0到4这5个行索引的所有可能排列，每个排列对应一种行的顺序。
• 批量索引：把排列转成NumPy数组后，直接用arr[perm_indices]就能一次性提取所有排列后的行组合——因为perm_indices的每一行是一组行索引，NumPy会自动按组提取对应的行，堆叠成三维数组。
• 行内元素保持不变：这个方法只是改变行的顺序，每行内部的元素顺序完全和原数组一致，符合需求。

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