是否有Numpy函数可按值归类多数组元素并提取索引用于排序？

我有4个二维Numpy数组arr1、arr2、arr3和arr4，每个数组仅包含1、2、3、4四个值，且同一位置的数组值互不重复（例如若arr1[i,j]=1，则其他数组该位置不会是1）。我需要生成按值归类的新数组，同时获取对应索引以用于其他数据的排序。目前我用多个if语句实现，但效率较低，想了解有没有合适的Numpy函数完成这个操作。

示例数组：

import numpy as np

arr1 = np.array([[1,2,3,4],[2,3,4,1]])
arr2 = np.array([[2,3,4,1],[3,4,1,2]])
arr3 = np.array([[3,4,1,2],[4,1,2,3]])
arr4 = np.array([[4,1,2,3],[1,2,3,4]])

还有4个存储测量数据的数组data1、data2、data3、data4：

data1 = np.array([[10,11,12,13],[14,15,16,17]])
data2 = np.array([[20,21,22,23],[24,25,26,27]])
data3 = np.array([[30,31,32,33],[34,35,36,37]])
data4 = np.array([[40,41,42,43],[44,45,46,47]])

预期输出：

new_data1 = np.array([[10,41,32,23],[44,35,26,17]])
new_data2 = np.array([[20,11,42,33],[14,45,36,27]])
new_data3 = np.array([[30,21,12,43],[24,15,46,37]])
new_data4 = np.array([[40,31,22,13],[34,25,16,47]])

可以利用Numpy的广播和高级索引实现高效处理，完全替代循环与if判断，步骤如下：

1. 将所有arr数组堆叠为三维数组，同时将data数组也做同样堆叠，方便统一操作；
2. 针对每个目标值（1-4），找出堆叠后arr数组中对应值的位置索引；
3. 利用索引从堆叠的data数组中提取对应数据，生成最终结果。

具体代码实现：

import numpy as np

# 输入数组
arr1 = np.array([[1,2,3,4],[2,3,4,1]])
arr2 = np.array([[2,3,4,1],[3,4,1,2]])
arr3 = np.array([[3,4,1,2],[4,1,2,3]])
arr4 = np.array([[4,1,2,3],[1,2,3,4]])

data1 = np.array([[10,11,12,13],[14,15,16,17]])
data2 = np.array([[20,21,22,23],[24,25,26,27]])
data3 = np.array([[30,31,32,33],[34,35,36,37]])
data4 = np.array([[40,41,42,43],[44,45,46,47]])

# 堆叠arr和data数组，生成三维数组
arr_stack = np.stack([arr1, arr2, arr3, arr4], axis=0)
data_stack = np.stack([data1, data2, data3, data4], axis=0)

# 生成每个目标值对应的new_data
new_data_list = []
for target in [1,2,3,4]:
    # 获取arr_stack中等于目标值的元素在第一个维度的索引，并恢复为原数组形状
    idx = np.where(arr_stack == target)[0].reshape(arr1.shape)
    # 生成行、列的广播索引，结合idx提取对应data数据
    row_idx = np.arange(arr1.shape[0])[:, None]
    col_idx = np.arange(arr1.shape[1])
    new_data = data_stack[idx, row_idx, col_idx]
    new_data_list.append(new_data)

new_data1, new_data2, new_data3, new_data4 = new_data_list

# 验证输出
print("new_data1:\n", new_data1)
print("new_data2:\n", new_data2)
print("new_data3:\n", new_data3)
print("new_data4:\n", new_data4)

代码说明 #

• np.stack：将多个二维数组沿第一个维度堆叠为三维数组，实现批量统一处理；
• np.where(arr_stack == target)[0]：定位所有等于目标值的元素所属的原始arr数组索引，通过reshape恢复为与原数组一致的形状；
• 高级索引：通过生成行、列的广播索引，配合idx精准提取每个位置对应目标值的data数据，全程为向量化操作，效率远高于循环判断。

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