Numpy中array[:,*idx]的替代方案：可变维度多索引循环赋值

嘿，我之前也踩过这个坑！Python确实不允许在Numpy数组的索引方括号里直接用*idx这种解包语法，但咱们有几个非常实用的替代方案，完全能满足你需要遍历所有{0,1,...,m}^n索引组合的需求：

核心思路：用元组构造完整索引 #

Numpy的索引其实是支持元组的，所以我们只需要把你的idx列表转换成元组，再和前面的切片（比如:）拼接成一个完整的索引元组就行。

单条索引赋值的示例 #

import numpy as np

# 初始化4维数组
a = np.ones([3]*4)
idx = [2, 1, 0]

# 构造索引元组：第一个维度用slice(None)代替冒号，后面跟上转成tuple的idx
indices = (slice(None),) + tuple(idx)
# 执行赋值，效果和a[:,2,1,0] = 1完全一致
a[indices] = 1

# 验证一下结果
print(np.array_equal(a[:,2,1,0], np.ones(3)))  # 输出True

这里slice(None)就是冒号:的等价写法，用它是为了方便和tuple拼接，如果你有多个前置维度，也可以灵活扩展这个索引元组，比如(slice(None), slice(None)) + tuple(idx)对应a[:,:,...]的场景。

遍历所有索引组合的方案 #

要遍历{0,1,...,m}^n中的所有idx，可以用itertools.product生成所有笛卡尔积组合，它直接返回tuple类型的索引，连转换都省了：

import numpy as np
import itertools

m = 2
n = 3
# 初始化n+1维数组（示例中是4维）
a = np.ones([3]*(n+1))

# 遍历所有可能的idx组合：itertools.product生成所有笛卡尔积
for idx in itertools.product(range(m+1), repeat=n):
    # 直接拼接索引元组，idx已经是tuple了
    indices = (slice(None),) + idx
    # 这里可以根据需求赋值，比如随idx变化的不同值
    a[indices] = 1

为什么a[:,*idx]不行？ #

简单来说就是Python的语法限制——索引方括号内只能用逗号分隔的表达式，不支持*解包操作。而Numpy接受tuple作为索引，所以把列表转成tuple是最直接的兼容方式。

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