New list based on indices from another list in Python

I have an array X and a list A1. I want to create a new list B1 such that it consists of values from X corresponding to indices in A1. For example, the code should pick values from X[0] for indices in A1[0] and so on…I present the current and expected outputs.

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)

The current output is

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

The expected output is

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

Answer

For the ith element of A1, you want to pick elements from the ith row of X, so iterate over A1 and X simultaneously using zip. In the loop, ai is a list containing a single list. This inner list contains the indices you want.

result = []
for xi, ai in zip(X, A1):
    indices = ai[0]
    result.append(xi[indices].tolist())

Which gives the desired result:

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

Note that I converted xi[indices] to a list, but if the result you want is a numpy array then you don’t need to do that, and instead just:

result = np.array([xi[ai[0]] for xi, ai in zip(X, A1)])

which gives a 2x3 result array:

array([[353.856161,   0.      ,   0.      ],
       [  0.      , 294.983702, 126.991664]])