任意维度嵌套列表如何实现类似NumPy数组的元组索引与切片赋值？

不需要递归实现，直接通过逐层迭代定位到目标层级即可，同时天然支持整数索引、切片对象，适配任意维度场景。

核心实现 #

取值功能 #

def get_nd(ndlist, indexes):
    current = ndlist
    for idx in indexes:
        current = current[idx]
    return current

赋值功能 #

赋值时只需遍历到倒数第二层的父列表，再用最后一个索引完成赋值即可：

def set_nd(ndlist, indexes, value):
    current = ndlist
    for idx in indexes[:-1]:
        current = current[idx]
    current[indexes[-1]] = value

使用示例 #

普通索引赋值 #

以3维2x2x2列表为例：

ndlist = [[[10, 10], [10, 10]], [[10, 10], [10, 10]]]
# 对应(0,0,0)位置赋值为11
set_nd(ndlist, (0,0,0), 11)
print(ndlist[0][0][0]) # 输出11

切片赋值 #

如果要实现类似NumPy的ndarray[0:2, 0, 0] = [0, 0]效果，直接将切片对象放入索引元组即可：

# 切片对象slice(0,2)对应Python语法里的0:2
set_nd(ndlist, (slice(0,2), 0, 0), [0, 0])
print(ndlist[0][0][0]) # 输出0
print(ndlist[1][0][0]) # 输出0

类封装实现类NumPy索引语法 #

如果想要和NumPy数组一样直接通过ndlist[indexes]的方式调用，可以自定义包装类重写魔术方法：

class NDList:
    def __init__(self, data):
        self.data = data

    def __getitem__(self, key):
        # 兼容单个索引和元组索引
        keys = key if isinstance(key, tuple) else (key,)
        current = self.data
        for k in keys:
            current = current[k]
        return NDList(current) if isinstance(current, list) else current

    def __setitem__(self, key, value):
        keys = key if isinstance(key, tuple) else (key,)
        current = self.data
        for k in keys[:-1]:
            current = current[k]
        current[keys[-1]] = value

使用示例：

ndl = NDList([[[10, 10], [10, 10]], [[10, 10], [10, 10]]])
# 直接用元组索引赋值
ndl[(0,0,0)] = 11
# 切片赋值
ndl[(slice(0,2), 0, 0)] = [0, 0]

方案优势 #

• 时间复杂度为O(ndim)，和递归效率一致，但没有递归深度过大导致栈溢出的风险，支持任意高维度场景
• 原生支持切片、整数索引，逻辑和NumPy索引完全对齐
• 实现轻量，无第三方依赖

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