Bera
Bera

Reputation: 1949

Numpy where behavior

Why is numpy.where returning two arrays when only one True value is found but when n True>1 return one array for each True. I would like it to always return one array for each true.

import numpy as np
a = np.zeros((5,5),dtype=int)
print(a)
#[[0 0 0 0 0]
# [0 0 0 0 0]
# [0 0 0 0 0]
# [0 0 0 0 0]
# [0 0 0 0 0]]

a[1,1]=1
print(np.where(a==1))
#(array([1]), array([1])) #Two separate arrays
#Why not array([1,1])


a[1,2]=1
print(np.where(a==1))
#(array([1, 1]), array([1, 2])) #One array for each True, desired behavior

Upvotes: 1

Views: 146

Answers (2)

Madhoolika
Madhoolika

Reputation: 444

np.where returns the row and column indices where the condition specified holds True. In the example your considered, coincidentally the row column indices directly matched the positions where the condition is satisfied. Hence take one more location as 1 to understand np.where.

import numpy as np
a = np.zeros((5,5),dtype=int)
a[1,1]=1
a[1,2]=1
a[2,3] = 1

row_ind, col_ind = np.where(a==1)

row_ind
Out[22]: array([1, 1, 2], dtype=int64)

col_ind
Out[23]: array([1, 2, 3], dtype=int64)

[(i,j) for i, j in zip(row_ind, col_ind)]
Out[24]: [(1, 1), (1, 2), (2, 3)]

Upvotes: 1

amzon-ex
amzon-ex

Reputation: 1744

The way you interpret this is not quite correct. For a 2-D array, numpy.where returns two sets of indices:the first array contains indices along the first axis (rows), the second array the indices along the second axis (columns). (And so on - for a 3-D array, you'll get three arrays, one for the indices of each axis).

So, you should really read (array([1,2]), array([1,3])) as follows: the value 1 is present in the locations (1, 1) and (2, 3).
The reason for such a return is, if you mask the original array with the returned tuple of arrays, you get only the values at the locations where your condition evaluated to true.

Upvotes: 1

Related Questions