如何用不等长列索引列表切片数组并批量计算行向量范数？

我有一个N×M的numpy数组，以及一组用于切片数组的任意列索引集合列表。例如，3×3数组：

my_arr = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])

索引集合为：

my_idxs = [[0, 1], [2]]

我希望使用这些索引集合选择对应列，并通过np.linalg.norm()计算每行向量的范数，且对所有索引集合执行该操作。上述示例应得到：

[[2.23606797749979, 3],
 [2.23606797749979, 3],
 [2.23606797749979, 3]]

当所有索引集合长度相同时（如my_idxs = [[0, 1], [1, 2]]），可直接用np.linalg.norm(my_arr[:, my_idxs], axis=1)得到结果：

[[2.23606797749979, 3.605551275463989],
 [2.23606797749979, 3.605551275463989],
 [2.23606797749979, 3.605551275463989]]

但当索引集合长度不同时（如my_idxs = [[0, 1], [2]]），不规则的索引形状会导致切片报错。请问能否实现无需循环处理单个索引集合的单行解决方案？

可以通过numpy的向量化操作实现无循环的高效解决方案，以下是两种实用方法：

方法1：掩码广播法 #

先构建标记索引集合的掩码矩阵，再通过广播计算每行对应列的范数：

import numpy as np

my_arr = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
my_idxs = [[0, 1], [2]]

# 创建对应索引集合的掩码矩阵，形状为(索引集合数量, 数组列数)
mask = np.zeros((len(my_idxs), my_arr.shape[1]), dtype=bool)
for i, idxs in enumerate(my_idxs):
    mask[i, idxs] = True

# 广播计算范数：先对元素平方，按掩码筛选后求和，最后开平方
result = np.sqrt((my_arr[:, None, :] ** 2 * mask).sum(axis=2))

执行结果：

array([[2.23606798, 3.        ],
       [2.23606798, 3.        ],
       [2.23606798, 3.        ]])

方法2：扁平化索引分组求和法 #

将所有索引扁平化后，利用分组统计计算平方和，再开平方得到范数：

import numpy as np

my_arr = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
my_idxs = [[0, 1], [2]]

# 扁平化所有索引，并生成对应分组标记
flat_idxs = np.concatenate(my_idxs)
groups = np.repeat(np.arange(len(my_idxs)), [len(x) for x in my_idxs])

# 按分组统计每行元素的平方和，再开平方
sums = np.bincount(groups, weights=my_arr[:, flat_idxs] ** 2, minlength=len(my_idxs)).reshape(my_arr.shape[0], -1)
result = np.sqrt(sums)

这种方法避免了显式构建掩码，适合索引集合数量多、长度差异大的场景。

说明 #

两种方法都没有循环处理单个索引集合，完全依赖numpy的向量化操作实现，既能保证代码简洁性，又能维持高效的计算性能。

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