Using list comprehension to search a 2d array (python)

Question

I am trying to search a 2D array of characters and return the array indices, (x_T,y_T), of all of the letter T's in the array. I figure this could easily done with two stacked for loops but I was curious to know if it could be done my efficiently using list comprehension. I have tried things like:

TPos = [[x, y] for x in range(0, len(array[x])) for y in range(0, len(array[x][y])) if array[x][y] == 'T']

But I am getting errors to do with array bounds. Could someone point me in the right direction. Cheers, Jack

Edit

Now trying to use 'ndenumerate' like so:

TPos = [pos for pos, x in numpy.ndenumerate(array) if x == "T"]

Cory Kramer · Accepted Answer · 2014-09-26T17:19:44.640

3

This ended up being a beast to one-line, for readability it may be easier to break it up.

l = [['a','b','T'],
     ['T','T','g'],
     ['c','T','T']]

reduce(lambda x,y: x+y, [[(x,y) for y in range(len(l[y])) if l[x][y] == 'T'] for x in range(len(l))])

Output

[(0, 2), (1, 0), (1, 1), (2, 1), (2, 2)]

Edit With numpy this is much simpler.

import numpy as np
[pos for pos, x in np.ndenumerate(np.array(l)) if x == 'T']

Output

[(0, 2), (1, 0), (1, 1), (2, 1), (2, 2)]

edited Sep 26 '14 at 17:19

answered Sep 26 '14 at 17:14

Cory Kramer

114,268
16
167
218

Cheers Cyber, cool answer. Any idea why: TPos = [pos for pos, x in numpy.ndenumerate(array) if x == "T"] doesnt work? – JMzance Sep 26 '14 at 17:16
I didn't even know about `numpy.ndenumerate`, that is much simpler! – Cory Kramer Sep 26 '14 at 17:18
yeah its great! Cant't get it to work though. Returns empty array – JMzance Sep 26 '14 at 17:19
I just tried it with my sample array `l` from above, it worked just as expected. – Cory Kramer Sep 26 '14 at 17:20
Ah yeah I was being daft. ndeumerate wtf – JMzance Sep 26 '14 at 18:23
Can you do the same thing for y[pos for pos, x in np.ndenumerate(np.array(l)) if x == 'T'] – Pat Dec 06 '14 at 22:03

score 1 · Answer 2 · answered Sep 27 '14 at 03:10

numpy where does this simply:

In [16]: l = [['a','b','T'],
   ....:      ['T','T','g'],
   ....:      ['c','T','T']]

In [20]: np.where(np.array(l)=='T')
Out[20]: (array([0, 1, 1, 2, 2]), array([2, 0, 1, 1, 2]))

In [22]: zip(*np.where(np.array(l)=='T'))
Out[22]: [(0, 2), (1, 0), (1, 1), (2, 1), (2, 2)]

Without numpy here's my version of the list comprehension:

In [27]: [(i,j) for i,v in enumerate(l) for j,k in enumerate(v) if k=='T']
Out[27]: [(0, 2), (1, 0), (1, 1), (2, 1), (2, 2)]

nice inversion of the iteration to go straight to the list of tuples -- thanks -- I learned something! — jstevenco, Sep 28 '14 at 18:30

Sarah · Answer 3 · 2014-09-26T17:37:57.490

0

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

def return_indexes(some_list,what_you_want_to_find):
    temp=[]
    for i in range(len(some_list)):
        for j in range(len(a[i])):
            if a[i][j] == what_you_want_to_find:
                temp.append([i,j])
    return temp

print return_indexes(a,7)

edited Sep 26 '14 at 17:37

answered Sep 26 '14 at 17:25

Sarah

11
2

2

This is a viable function, but the question is asking specifically for list comprehensions. – MackM Sep 26 '14 at 17:26
This function is not correct. It will return only the first match, not all matches. – Cory Kramer Sep 26 '14 at 17:30
I didn't see Cyber's awesome answer until after my post. – Sarah Sep 26 '14 at 17:31

score 0 · Answer 4 · edited May 23 '17 at 11:46

Not quite as nice as @Cyber's answer (i.e., list of tuple lists rather than list of tuples), but uses "vanilla" list comprehensions with enumerate (hey, I got curious):

l = [['a','b','T'],['T','T','g'],['c','T','T']]
[[(i,j) for j,col in enumerate(row) if col=='T'] for i,row in enumerate(l)]

Output

[[(0, 2)], [(1, 0), (1, 1)], [(2, 1), (2, 2)]]

of course you could apply reduce(lambda ...) or similar approach to this to get a list of tuples.

Using list comprehension to search a 2d array (python)

Edit

4 Answers4