Jimmy
Jimmy

Reputation: 1

Return which elements in two numpy arrays are both equal to X value?

For example, A = [5, 0, 0, 5] B = [5, 5, 0, 5]

I want to get which elements are equal to 5 in both arrays

The output should be = [0, 3]

Upvotes: 0

Views: 72

Answers (2)

mathfux
mathfux

Reputation: 5949

In one dimensional arrays np.flatnonzero works perfectly:

import numpy as np
A = np.array([5, 0, 0, 5])
B = np.array([5, 5, 0, 5])
np.flatnonzero((A == 5) & (B == 5))

np.where or np.nonzero, however, returns tuple and you need to take its first item:

np.where((A == 5) & (B == 5))[0] or

np.nonzero((A == 5) & (B == 5))[0]

Output

array([0, 3], dtype=int64)

Remark: make sure you converted A and B to arrays first.

Upvotes: 0

Marcin
Marcin

Reputation: 1381

np.where(np.logical_and(A == 5, B == 5))

Upvotes: 1

Related Questions