如何高效切片大型二维NumPy数组并按索引分组？

你的代码慢的核心问题是遍历了131000行的索引值，而不是仅遍历0-25这26个唯一索引，导致重复执行了上万次相同的掩码计算和切片操作，完全没必要。

下面是两种更高效的方案：

方案一：仅遍历26个唯一索引（简单易实现） #

直接循环0到25这26个值，每次计算对应索引的掩码并切片，比原代码快几个数量级：

import numpy as np

# 提取第一列的索引行
idx = huge_array[0]

# 仅循环26个唯一索引
for index in range(26):
    # 一次计算掩码，获取对应列的索引
    mask = idx == index
    # 切片得到对应数组
    new_array = huge_array[:, mask]
    # 保存数组
    np.save(f"array_{index}.npy", new_array)

方案二：排序后分割（性能最优） #

如果数据量极大，排序后分割的性能会更优，因为排序后相同索引的列连续，分割操作是O(n)时间复杂度：

import numpy as np

# 提取第一列的索引行
idx = huge_array[0]
# 获取排序后的列索引
sorted_col_indices = np.argsort(idx)
# 对原数组按索引排序，相同索引的列会连续排列
sorted_array = huge_array[:, sorted_col_indices]
# 获取每个索引在排序后数组中的起始位置
_, split_positions = np.unique(idx[sorted_col_indices], return_index=True)
# 分割数组，得到26个对应索引的子数组
split_arrays = np.split(sorted_array, split_positions[1:], axis=1)

# 遍历保存每个子数组
for index, arr in enumerate(split_arrays):
    np.save(f"array_{index}.npy", arr)

为什么这两种方法更快？ #

• 方案一把循环次数从131000次降到26次，避免了大量重复计算；
• 方案二通过一次排序将相同索引的列聚合，分割操作无需重复遍历整个数组，适合超大规模数据。

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