CODEWITHSUNDEEP
pythonnumpy
lemon
lemon

Reputation: 747

Numpy indexing using an array of booleans

I have the following array which indicates whether or not to take a certain item:

import numpy as np

test_array = np.array([[0, 0, 1],
                       [1, 1, 0],
                       [1, 1, 1]])

The array I want to index is this one:

classes = ['a', 'b', 'c']

This is what the result should be:

[['c'], ['a', 'b'], ['a', 'b', 'c']]

How can this be done?

Upvotes: 1

Views: 110

Answers (6)

AMC
AMC

Reputation: 2702

The answers I've seen so far range from awkward to, quite frankly, baffling, so here is a straightforward solution.

import np

arr = np.array([[0, 0, 1], [1, 1, 0], [1, 1, 1]])

arr_bools = arr.astype(np.bool)

lookup_lst = np.array(['a', 'b', 'c'])

res = [lookup_lst[row].tolist() for row in arr_bools]

Upvotes: 1

Paul Panzer
Paul Panzer

Reputation: 53109

This can be done by matrix multiplication: 

[*map(list, test_array.astype('O')@classes)]
# [['c'], ['a', 'b'], ['a', 'b', 'c']]

Upvotes: 1

xenteros
xenteros

Reputation: 15862

I would start with something like that:

result = []
for row in test_array:
    partial_result = []
    for i in range(3):
        if row[i] == 1:
            partial_result.append(classes[i])
    result.append(partial_result)
print(result)

Results with:

[['c'], ['a', 'b'], ['a', 'b', 'c']]

In Python, we prefer list comprehension over loops, so time to improve:

print([[classes[i] for i, val in enumerate(row) if val] for row in test_array])

enumerate is an in-build function which takes an iterable object as a parameter and returns iterable of tuples (index, element) for all elements in the original iterable, so enumerate(row) will return (0, [0, 0, 1]), (1, [1, 1, 0]) and (2, [1, 1, 1]).

for i, val in enumerate(row) if val will work, because 1s are interpreted as True in Python and 0s as False.

[[classes[i] for i, val in enumerate(row) if val] for row in test_array]
^ create a list of elements based on some original list ------->^
 ^ each element of that list will be a list itself.
      ^ elements of that inner lists will be objects from classes list
              ^ for each pair (i, element) from enumerate(row) take this ith
                element, but just if val == 1 ^

Upvotes: 1

Reznik
Reznik

Reputation: 2806

I would do it as this:

result = []
for array in test_array:
     result.append([classes[i] for i,value in enumerate(array ) if value ])

Upvotes: 1

SPH
SPH

Reputation: 510

In one line you can do: 

print ([[c for (x, c) in zip(l, classes) if x] for l in test_array])

Upvotes: 1

Dani Mesejo
Dani Mesejo

Reputation: 61930

You could do the following:

import numpy as np

test_array = np.array([[0, 0, 1],
                       [1, 1, 0],
                       [1, 1, 1]])

classes = ['a', 'b', 'c']

lookup = dict(enumerate(classes))
result = [[lookup[i] for i, e in enumerate(arr) if e] for arr in test_array]
print(result)

Output

[['c'], ['a', 'b'], ['a', 'b', 'c']]

Upvotes: 1

Related Questions