如何按索引筛选字典中不同形状的numpy数组？

这个问题我之前也碰到过！核心痛点在于不同形状的数组，索引的维度需要对应调整——比如一维数组直接用v[idx]就够，但二维数组（比如(3,1000)）得在最后一维取索引，也就是v[:, idx]，更高维的数组同理。

我们可以写一个通用的筛选逻辑，自动适配任意维度的数组，只要它们的最后一维长度和索引范围匹配就行：

import numpy as np
import random

# 生成不同形状的测试数据
a1 = np.random.uniform(0., 1., 1000)  # 一维数组 (1000,)
a2 = np.random.uniform(0., 1., (3, 1000))  # 二维数组 (3,1000)
a3 = np.random.uniform(0., 1., (5, 2, 1000))  # 三维数组 (5,2,1000)

# 随机选取目标索引
idx = random.sample(range(0, 1000), 25)

# 构建原始字典
aa = {'a1': a1, 'a2': a2, 'a3': a3}

# 通用筛选逻辑：针对最后一维应用索引
bb = {}
for k, v in aa.items():
    # 根据数组维度生成切片元组：前面所有维度取全部，最后一维用目标索引
    slice_tuple = tuple([slice(None)] * (v.ndim - 1) + [idx])
    bb[k] = v[slice_tuple]

# 验证结果形状是否符合预期
print(bb['a1'].shape)  # 输出 (25,)
print(bb['a2'].shape)  # 输出 (3,25)
print(bb['a3'].shape)  # 输出 (5,2,25)

关键思路拆解： #

• 用v.ndim获取数组的维度数，比如一维数组返回1，二维返回2
• 生成切片元组：前面的所有维度用slice(None)（等价于冒号:，表示取该维度的全部元素），最后一维替换为我们的目标索引idx
• 这种写法完全通用，不管数组是1维、2维还是更高维，只要最后一维长度和索引匹配，就能正确筛选

如果想更简洁，也可以用字典推导式一行完成：

bb = {k: v[tuple([slice(None)]*(v.ndim-1)+[idx])] for k, v in aa.items()}

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