Finding the index of elements based on a condition using python list comprehension

Question

The following Python code appears to be very long winded when coming from a Matlab background

>>> a = [1, 2, 3, 1, 2, 3] >>> [index for index,value in enumerate(a) if value > 2] [2, 5] 

When in Matlab I can write:

>> a = [1, 2, 3, 1, 2, 3]; >> find(a>2) ans =      3     6 

Is there a short hand method of writing this in Python, or do I just stick with the long version?

---

Thank you for all the suggestions and explanation of the rationale for Python's syntax.

After finding the following on the numpy website, I think I have found a solution I like:

http://docs.scipy.org/doc/numpy/user/basics.indexing.html#boolean-or-mask-index-arrays

Applying the information from that website to my problem above, would give the following:

>>> from numpy import array >>> a = array([1, 2, 3, 1, 2, 3]) >>> b = a>2  array([False, False, True, False, False, True], dtype=bool) >>> r = array(range(len(b))) >>> r(b) [2, 5] 

The following should then work (but I haven't got a Python interpreter on hand to test it):

class my_array(numpy.array):     def find(self, b):         r = array(range(len(b)))         return r(b)   >>> a = my_array([1, 2, 3, 1, 2, 3]) >>> a.find(a>2) [2, 5] 

Answer 1

Another way:

>>> [i for i in range(len(a)) if a[i] > 2] [2, 5] 

In general, remember that while find is a ready-cooked function, list comprehensions are a general, and thus very powerful solution. Nothing prevents you from writing a find function in Python and use it later as you wish. I.e.:

>>> def find_indices(lst, condition): ...   return [i for i, elem in enumerate(lst) if condition(elem)] ...  >>> find_indices(a, lambda e: e > 2) [2, 5] 

Note that I'm using lists here to mimic Matlab. It would be more Pythonic to use generators and iterators.

Answer 2

• In Python, you wouldn't use indexes for this at all, but just deal with the values—[value for value in a if value > 2]. Usually dealing with indexes means you're not doing something the best way.

• If you do need an API similar to Matlab's, you would use numpy, a package for multidimensional arrays and numerical math in Python which is heavily inspired by Matlab. You would be using a numpy array instead of a list.

  >>> import numpy >>> a = numpy.array([1, 2, 3, 1, 2, 3]) >>> a array([1, 2, 3, 1, 2, 3]) >>> numpy.where(a > 2) (array([2, 5]),) >>> a > 2 array([False, False,  True, False, False,  True], dtype=bool) >>> a[numpy.where(a > 2)] array([3, 3]) >>> a[a > 2] array([3, 3])