CODEWITHSUNDEEP
pythonnumpy
Steven G
Steven G

Reputation: 17152

Find indices of the elements smaller than x in a numpy array

Assuming that I have a numpy array such as:

import numpy as np    
arr = np.array([10,1,2,5,6,2,3,8])

How could I extract an array containing the indices of the elements smaller than 6 so I get the following result:

np.array([1,2,3,5,6])

I would like something that behave like np.nonzero() but instead of testing for nonzero value, it test for value smaller than x

Upvotes: 8

Views: 34614

Answers (3)

Giovanni Cavallin
Giovanni Cavallin

Reputation: 159

I'd suggest a cleaner and self-explainable way to do so: First, find the indices where the condition is valid:

>> indices = arr < 6
>> indices
>> [False, True, True, True, False, True, False]

Then, use the indices for indexing:

>> arr[indices]
>> [1, 2, 5, 2, 3]

or for finding the right position in the original array:

>> np.where(indices)[0]
>> [1, 2, 3, 5, 6]

Upvotes: 2

Raikan 10
Raikan 10

Reputation: 115

The simplest way one can do this is by

arr[arr<6]

Upvotes: 4

akuiper
akuiper

Reputation: 215137

You can use numpy.flatnonzero on the boolean mask and Return indices that are non-zero in the flattened version of a:

np.flatnonzero(arr < 6)
# array([1, 2, 3, 5, 6])

Another option on 1d array is numpy.where:

np.where(arr < 6)[0]
# array([1, 2, 3, 5, 6])

Upvotes: 12

Related Questions