如何高效为嵌套二维列表的指定行批量追加元素以优化耗时？

需求与现状 #

我有如下嵌套列表：

a = [[[0, 1], [1, 1]], [[1, 0], [1, 1]]]

它包含两个二维子列表：

a[0] = [[0, 1], [1, 1]]           
a[1] = [[1, 0], [1, 1]]

需要实现给每个二维子列表的第一行追加0、第二行追加1，最终得到：

a[0] = [[0, 1, 0], [1, 1, 1]]           
a[1] = [[1, 0, 0], [1, 1, 1]]

当前通过嵌套循环实现，但矩阵规模扩大后耗时显著增加，代码如下：

elif (rown == 0) and (columnn > 0):
    s_routine = time.time()
    p_temp = p[:]
    p = []
    for matrix in p_temp:
        ca = matrix[0][columnn]
        cb = matrix[1][columnn]
        if columnv == 0:
            zl = zfc[str(ca)+str(cb)]
            for z in zl:
                a = matrix[0] + [z[0][1]]
                b = matrix[1] + [z[1][1]]
                c = [a, b]
                p.append(c)
        elif columnv == 1:
            ol = zfc[str(ca)+str(cb)]
            for o in ol:
                a = matrix[0] + [o[0][1]]
                b = matrix[1] + [o[1][1]]
                c = [a, b]
                p.append(c)
    e_routine = time.time() - s_routine

优化方案 #

1. 用列表推导式替代嵌套循环 #

Python的列表推导式基于C实现，比手动嵌套循环效率更高，可以简化内层循环逻辑：

elif (rown == 0) and (columnn > 0):
    s_routine = time.time()
    p = []
    # 根据columnv选择对应的映射列表
    target_map = zfc if columnv == 0 else zfc
    for matrix in p[:]:
        ca = matrix[0][columnn]
        cb = matrix[1][columnn]
        item_list = target_map[f"{ca}{cb}"]
        # 用列表推导式生成所有新矩阵并批量追加
        p.extend([
            [matrix[0] + [item[0][1]], matrix[1] + [item[1][1]]]
            for item in item_list
        ])
    e_routine = time.time() - s_routine

如果只是针对示例中的a列表做批量追加，还可以用更简洁的嵌套列表推导式：

result = [[row + [0] if idx == 0 else row + [1] for idx, row in enumerate(mat)] for mat in a]

2. 用NumPy向量化操作处理大规模数据 #

当矩阵规模很大时，推荐使用NumPy的向量化操作，完全避开Python层面的循环，利用底层计算能力：

import numpy as np

# 将原列表转为NumPy数组
a_np = np.array(a)
# 构造要追加的列：每个二维矩阵的第一行加0，第二行加1
append_col = np.array([[0], [1]])
# 对每个二维矩阵批量拼接列
result_np = np.concatenate([a_np, np.tile(append_col, (a_np.shape[0], 1, 1))], axis=2)
# 如需转回Python列表
result = result_np.tolist()

这种方式在数据量越大时，性能提升越明显，适合大规模矩阵处理场景。

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