Python Numpy 2-dimensional array iteration

Question

Would like to build a list of indices into a 2 dimensional bool_ array, where True.

import numpy
arr = numpy.zeros((6,6), numpy.bool_)
arr[2,3] = True
arr[5,1] = True
results1 = [[(x,y) for (y,cell) in enumerate(arr[x].flat) if cell] for x in xrange(6)]
results2 = [(x,y) for (y,cell) in enumerate(arr[x].flat) if cell for x in xrange(6)]

results 1:

[[], [], [(2, 3)], [], [], [(5, 1)]]

results 2 is completely wrong

Goal:

[(2, 3),(5, 1)]

Any way to do this without flattening the list afterwards, or any better way to do this in general?

jterrace · Accepted Answer · 2011-11-14 17:46:23Z

1

I think the function you're looking for is numpy.where. Here's an example:

>>> import numpy
>>> arr = numpy.zeros((6,6), numpy.bool_)
>>> arr[2,3] = True
>>> arr[5,1] = True
>>> numpy.where(arr)
(array([2, 5]), array([3, 1]))

You can turn this back into an index like this:

>>> numpy.array(numpy.where(arr)).T
array([[2, 3],
       [5, 1]])

answered Nov 14, 2011 at 17:46

jterrace

67.4k23 gold badges164 silver badges208 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

user1012037 Over a year ago

Oh dear, hadn't heard of that. zip(*numpy.where(arr)) is working nicely. I'll leave this open for awhile to hear if anyone else has alternatives.

jfs Over a year ago

np.where() with a single argument is equivalent to np.nonzero(). To transform to the OP's format: np.transpose(np.nonzero(a)) that is equivalent to np.argwhere(a).

jfs · Accepted Answer · 2011-11-14 17:59:55Z

0

>>> import numpy as np
>>> arr = np.zeros((6,6), np.bool_)
>>> arr[2,3] = True
>>> arr[5,1] = True
>>> np.argwhere(arr)
array([[2, 3],
       [5, 1]])

answered Nov 14, 2011 at 17:59

jfs

417k210 gold badges1k silver badges1.7k bronze badges

Collectives™ on Stack Overflow

Python Numpy 2-dimensional array iteration

2 Answers 2

2 Comments

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Related