How to check if nth row contains an nth element in numpy?

user_15001799 0
Viewed 49

Let's say I have a 2-d array like:

a = np.array([[0 0 1], [0 2 0], [3, 1 3]])
targets = np.array([1, 2, 3])

I want to find all the locations of 1 in the 1st row, 2 in the 2nd row, and 3 in the 3rd row. Is there a clean way to do this in numpy?

result = [[False False True], [False True False], [True False True]]

So far, I've come up with this:

np.array([row == np.max(row) for row in test_array])

While this solution works, it'd be nicer if I could rely on parallelized library code to do this task.

numpy python
1 Answers

You can do

Method 1

import numpy as np

a = np.array([[0, 0, 1], [0, 2, 0], [3, 1, 3]])

b = np.max(a, axis=1).reshape((a.shape[0], 1)).repeat(a.shape[1], axis = 1) == a

print(b)

Output

[[False False  True]
 [False  True False]
 [ True False  True]]

Method 2 Another, slower, way is

import numpy as np

a = np.array([[0, 0, 1], [0, 2, 0], [3, 1, 3]])

for row, i in zip(a, np.max(a, axis = 1)):
    for idx, item in enumerate(row):
        if i != item:
            row[idx] = False
        else:
            row[idx] = True

a = np.array(a, dtype = bool)
print(a)

Comparisons All comparisons are based on creating a and creating the desired result. This shows time metrics for Method 1, Method 2, your method, and the first solution in the accepted answer in the suggested duplicate of this question.

# Method 1 
6.65 µs ± 44.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
# Method 2 
10.5 µs ± 44.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
# Your Method 
15.3 µs ± 149 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
# a.max(axis = 1, keepdims = 1) == a (from suggested duplicate)
4.24 µs ± 22.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

Old Answer Originally each row in the question had either a zero or a nonzero number so this applies for that case.

import numpy as np

a = (np.array([[0, 0, 1], [0, 2, 0], [3, 0, 3]]) != 0)

print(a)

Output

[[False False  True]
 [False  True False]
 [ True False  True]]
user_185544560
Related