Python中基于另一列表索引生成新列表的问题

我有一个numpy数组X和一个嵌套列表A1，想要生成新列表B1，其中B1[i]包含X[i]中对应A1[i]内索引的元素——也就是A1[0]里的索引选取X[0]的元素，A1[1]里的索引选取X[1]的元素。

以下是我的代码、当前输出及期望输出：

我的代码 #

import numpy as np

X= np.array([[417.551036, 0.0, 0.0, 353.856161, 0.0, 282.754301, 0.0, 0.0,
        134.119055, 63.4573886, 208.344718, 1e-24],
       [417.551036, 0.0, 332.821605, 294.983702, 0.0, 278.809292,
        126.991664, 0.0, 136.02651, 83.1512525, 207.329562, 1e-24]])

A1=[[[3, 4, 6]], [[1, 3, 6]]]

for i in range(0,len(A1)):
    for j in range(0,len(X)):
        B1 = [[X[j][i] for i in indices] for indices in A1[i]]
    print(B1)

当前输出 #

[[294.983702, 0.0, 126.991664]]
[[0.0, 294.983702, 126.991664]]

期望输出 #

[[353.856161, 0.0, 0.0]]
[[0.0, 294.983702, 126.991664]]

解决方法 #

原代码的问题在于嵌套循环逻辑错误，且索引引用混乱。我们需要将X的每一行与A1中对应的索引组直接配对，而不是遍历所有行。

修正后的代码 #

import numpy as np

X= np.array([[417.551036, 0.0, 0.0, 353.856161, 0.0, 282.754301, 0.0, 0.0,
        134.119055, 63.4573886, 208.344718, 1e-24],
       [417.551036, 0.0, 332.821605, 294.983702, 0.0, 278.809292,
        126.991664, 0.0, 136.02651, 83.1512525, 207.329562, 1e-24]])

A1=[[[3, 4, 6]], [[1, 3, 6]]]

# 用zip配对X的行和A1的索引组
B1 = [[[x_row[idx] for idx in indices[0]]] for x_row, indices in zip(X, A1)]

# 打印结果
for item in B1:
    print(item)

关键修正点 #

1. 移除多余循环：原代码中内层遍历X所有行的循环是不必要的，我们只需要A1[i]对应X[i]
2. 正确索引配对：通过zip(X, A1)直接将X的行和A1的索引组绑定，避免手动索引出错
3. 索引引用修正：针对A1的嵌套结构（每个元素是[[...]]），取出内部的索引列表indices[0]，再从对应X行中选取元素

输出结果 #

[[353.856161, 0.0, 0.0]]
[[0.0, 294.983702, 126.991664]]

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