New list based on indices from another list in Python-CodePudding

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]]

CodePudding user response：

for rowNumber, indices in enumerate(A1):
    B1 = [X[rowNumber][index] for index in indices]
    print(B1)

can be done with one list comprehension as well but little less readable IMO

list([X[rowNumber][index] for index in indices] for rowNumber, indices in enumerate(A1))

CodePudding user response：

You can zip the X and A1

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 = np.array([[[3, 4, 6]], [[1, 3, 6]]])
for x, a in zip(X, A1):
    print(x[a])

or use np.take_along_axis

A1 = A1.reshape((2, 3))
print(np.take_along_axis(X, A1, 1))

Output

[[353.856161   0.         0.      ]
 [  0.       294.983702 126.991664]]

CodePudding user response：

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]])

CodePudding user response：

This seems to work:

B1 = []
for i, row in enumerate(A1):
    values = []
    for inner_row in row:
        for index in inner_row:
            values.append(X[i][index])
    B1.append(values)

➜  arrays python main.py
[[353.856161, 0.0, 0.0], [0.0, 294.983702, 126.991664]]