Select values from list based on bools in numpy array-CodePudding

For each sublist in array b, return values from list a with same position as positive boolean in b sublist (i.e. where True).

import pandas as pd
import numpy as np

a = pd.Series([1, 3, 5, 7, 9])  # values to choose from
b = np.array([[False, True, False, True, False],  # based on bools
              [False, False, False, False, False]])

out = []
for i, v in enumerate(b):
    out.append([])
    for j in range(len(e)):
        if v[j]:
            out[i].append(a[j])

out = np.array(out)  # np.array([[3,7],[]])  # result

# In first sublist, True is on index 1 and 3 which corresponds to values 3 and 7.
# In second sublist, there is not True, hence empty.

The above seems too laborious and it is possibly not making use of numpy vectorization (it is slow on large data).

CodePudding user response：

You can simple use:

a2 = a.to_numpy()
out = [a2[x] for x in b]

output: [array([3, 7]), array([], dtype=int64)]

CodePudding user response：

alternatively just use array b as a mask eg:

out = a[b[0]].to_numpy()