寻找Numpy中唯一子数组的出现位置以加速矩阵填充

嘿，我来帮你搞定这个效率瓶颈问题！你的思路完全找对了——靠唯一子数组减少函数调用次数，绝对是处理这种大规模数组的最优方向，而用Numpy原生方法定位重复子数组的索引，可比手动算范数靠谱多了，还能避免范数相同但子数组不同的错误匹配。

下面是具体的实现步骤，都是Numpy原生的高效操作：

第一步：把B数组重塑成方便去重的形状 #

我们关心的是B中每个(2,)的子数组，所以先把(N,N,2)的B展平成(N*N, 2)的二维数组，这样每一行就对应原来的一个子数组：

B_reshaped = B.reshape(-1, 2)

第二步：获取唯一子数组和逆索引 #

用np.unique的axis=0参数按行（也就是按子数组）去重，同时通过return_inverse=True拿到每个原位置对应的唯一子数组索引：

unique_subarrays, inverse_indices = np.unique(B_reshaped, axis=0, return_inverse=True)

这里unique_subarrays是所有不重复的(2,)子数组，inverse_indices是长度为N*N的数组，每个元素表示原B_reshaped中对应行属于unique_subarrays里的第几个子数组。

第三步：获取每个唯一子数组在原B中的位置 #

如果需要知道每个唯一子数组具体出现在B的哪些(i,j)位置，可以先生成原B的索引网格，再按逆索引分组：

# 生成B的(i,j)索引网格
i_indices, j_indices = np.meshgrid(np.arange(N), np.arange(N), indexing='ij')
# 把索引展平成一维，和B_reshaped的行数对应
i_flat = i_indices.ravel()
j_flat = j_indices.ravel()

# 按逆索引排序，把相同子数组的索引聚在一起
sorted_idx = np.argsort(inverse_indices)
# 计算每个唯一子数组的分组分割点
split_points = np.cumsum(np.bincount(inverse_indices))[:-1]
# 分割得到每个唯一子数组对应的(i,j)索引组
grouped_i = np.split(i_flat[sorted_idx], split_points)
grouped_j = np.split(j_flat[sorted_idx], split_points)

现在grouped_i[k]和grouped_j[k]就对应第k个唯一子数组在B中所有的行、列索引。

第四步：批量处理并填充A #

接下来就可以只给每个唯一子数组调用一次average_lat_pos，然后批量填充到A的对应位置：

# 初始化A数组
A = np.zeros((2*N, 2*N))

for k, subarr in enumerate(unique_subarrays):
    # 仅调用一次函数
    result = average_lat_pos(subarr)
    # 获取当前子数组对应的所有(i,j)位置
    is_, js_ = grouped_i[k], grouped_j[k]
    # 这里根据你实际的填充规则调整，比如假设填充2i~2i+1行、2j~2j+1列的区域
    # 用向量化方式赋值比循环更快
    A[2*is_[:, None]:2*is_[:, None]+2, 2*js_[None, :]:2*js_[None, :]+2] = result

额外建议：尝试向量化average_lat_pos函数 #

如果average_lat_pos内部的逻辑可以改成接收(M,2)的数组（M是唯一子数组的数量），并返回对应长度的结果数组，那效率还能再上一个台阶：

# 假设函数支持向量化输入，直接传入所有唯一子数组
all_results = average_lat_pos(unique_subarrays)
# 把结果映射回原B的每个位置，再重塑成(N,N)的形状
result_map = all_results[inverse_indices].reshape(N, N)
# 如果你的填充规则是把每个result放到A的2i~2i+1、2j~2j+1区域，用kron直接生成A
A = np.kron(result_map, np.ones((2, 2)))

这种方式完全避免了循环，对于N=8000的规模来说，速度提升会非常明显。

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