如何通过维度索引列表访问与设置Python n维列表元素

问题场景 #

通过维度索引序列访问、修改多层嵌套Python列表时，可通过reduce(operator.getitem, index, lst)实现元素读取，但该写法无法直接用于赋值，直接执行reduce(operator.getitem, index, lst) = val会触发语法错误。
由于嵌套列表内元素数据类型不统一，无法使用numpy、torch等第三方库的高级索引能力，需要纯标准库实现方案。
问题复现代码如下：

import torch

tensor = torch.zeros(3, 5, 3)
tensor[2][1][0] = 1

lst = []
for i in range(3):
    lst.append([])
    for j in range(5):
        lst[i].append([])
        for k in range(3):
            lst[i][j].append([])

lst[2][1][0] = 1
index = [2, 1, 0]
tensor[[[i] for i in index]] = 2
print(tensor[2][1][0])

# 该行直接传入索引列表赋值会抛出错误
lst[[[i] for i in index]] = 2
print(lst[2][1][0])

实现方案 #

赋值操作的核心逻辑是：先定位到目标索引对应位置的父级容器，再取最后一维的索引对父容器执行赋值，避免直接取到元素值后赋值无法修改原列表的问题。
完整实现仅需调用Python标准库functools.reduce和operator模块：

from functools import reduce
import operator

def nested_get(nested_container, index_seq):
    """按索引序列读取嵌套结构中的元素"""
    return reduce(operator.getitem, index_seq, nested_container)

def nested_set(nested_container, index_seq, value):
    """按索引序列修改嵌套结构中的元素值"""
    # 定位到目标位置的上一级父容器
    parent = reduce(operator.getitem, index_seq[:-1], nested_container)
    # 用最后一个索引完成赋值
    parent[index_seq[-1]] = value

使用方法 #

# 构造测试嵌套列表
lst = [[[[] for _ in range(3)] for _ in range(5)] for _ in range(3)]
target_index = [2, 1, 0]

# 赋值
nested_set(lst, target_index, 2)
# 读取验证
print(nested_get(lst, target_index)) # 输出 2

补充说明 #

• 该方案兼容所有实现了__getitem__和__setitem__方法的嵌套结构，包括嵌套列表、嵌套字典、自定义索引对象，不局限于纯列表结构
• 如果使用的是类似[[2], [1], [0]]的花式多维索引格式，提前将索引展平为一维序列（如[2,1,0]）再传入函数即可

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