Find the indices of elements greater than x

Oliver Amundsen picture Oliver Amundsen · Dec 5, 2012 · Viewed 58.5k times · Source

Given the following vector,

a = [1, 2, 3, 4, 5, 6, 7, 8, 9]

I need to identify the indices of "a" whose elements are >= than 4, like this:

idx = [3, 4, 5, 6, 7, 8] 

The info in "idx" will be used to delete the elements from another list X (X has the same number of elements that "a"):

del X[idx] #idx is used to delete these elements in X. But so far isn't working.

I heard that numpy might help. Any ideas? Thanks!

Answer

Aesthete picture Aesthete · Dec 5, 2012
>>> [i for i,v in enumerate(a) if v > 4]
[4, 5, 6, 7, 8]

enumerate returns the index and value of each item in an array. So if the value v is greater than 4, include the index i in the new array.

Or you can just modify your list in place and exclude all values above 4.

>>> a[:] = [x for x in a if x<=4]
>>> a 
[1, 2, 3, 4]