Select elements of a numpy array based on the elements of a second array

Question

Consider a numpy array A of shape (7,6)

A = array([[0, 1, 2, 3, 5, 8],
           [4, 100, 6, 7, 8, 7],
           [8, 9,  10,  11, 5, 4],
           [12, 13, 14, 15, 1, 2],
           [1, 3, 5, 6, 4, 8],
           [12, 23, 12, 24, 4, 3],
           [1, 3, 5, 7,  89, 0]])

together with a second numpy array r of the same shape which contains the radius of A starting from a central point A(3,2)=0:

 r = array([[3, 3, 3, 3, 3, 4],
            [2, 2, 2, 2, 2, 3],
            [2, 1, 1, 1, 2, 3],
            [2, 1, 0, 1, 2, 3],
            [2, 1, 1, 1, 2, 3],
            [2, 2, 2, 2, 2, 3],
            [3, 3, 3, 3, 3, 4]])

I would like to pick up all the elements of A which are located at the position 1 of r, i.e. [9,10,11,15,4,6,5,13], all the elements of A located at position 2 of r and so on. I there some numpy function to do that? Thank you

Bi Rico · Accepted Answer · 2014-08-22 15:45:43Z

1

You can select a section of A by doing something like A[r == 1], to get all the sections as a list you could do [A[r == i] for i in range(r.max() + 1)]. This will work, but may be inefficient depending on how big the values in r go because you need to compute r == i for every i.

You could also use this trick, first sort A based on r, then simply split the sorted A array at the right places. That looks something like this:

r_flat = r.ravel()
order = r_flat.argsort()
A_sorted = A.ravel()[order]
r_sorted = r_flat[order]
edges = r_sorted.searchsorted(np.arange(r_sorted[-1] + 1), 'right')

sections = []
start = 0
for end in edges:
    sections.append(A_sorted[start:end])
    start = end

answered Aug 22, 2014 at 15:45

Bi Rico

25.9k3 gold badges57 silver badges75 bronze badges

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

1 Comment

jonrsharpe Over a year ago

It would probably be helpful to add the expected outputs for your various approaches.

jonrsharpe · Accepted Answer · 2014-08-22 15:44:30Z

0

I get a different answer to the one you were expecting (3 not 4 from the 4th row) and the order is slightly different (strictly row then column), but:

>>> A
array([[  0,   1,   2,   3,   5,   8],
       [  4, 100,   6,   7,   8,   7],
       [  8,   9,  10,  11,   5,   4],
       [ 12,  13,  14,  15,   1,   2],
       [  1,   3,   5,   6,   4,   8],
       [ 12,  23,  12,  24,   4,   3],
       [  1,   3,   5,   7,  89,   0]])
>>> r
array([[3, 3, 3, 3, 3, 4],
       [2, 2, 2, 2, 2, 3],
       [2, 1, 1, 1, 2, 3],
       [2, 1, 0, 1, 2, 3],
       [2, 1, 1, 1, 2, 3],
       [2, 2, 2, 2, 2, 3],
       [3, 3, 3, 3, 3, 4]])
>>> A[r==1]
array([ 9, 10, 11, 13, 15,  3,  5,  6])

Alternatively, you can get column then row ordering by transposing both arrays:

>>> A.T[r.T==1]
array([ 9, 13,  3, 10,  5, 11, 15,  6])

edited Aug 22, 2014 at 15:44

answered Aug 22, 2014 at 15:33

jonrsharpe

123k31 gold badges277 silver badges488 bronze badges

3 Comments

diegus Over a year ago

Thanks a lot.. That's how it should be, for r==1 it is 3 in place of 4!

diegus Over a year ago

For bigger array, Is it possible to loop over it in order to get more output arrays, i.e. for r=1,r=2,r=3...and so on?

jonrsharpe Over a year ago

@diegus yes, of course, you can loop over the values and create A[r==n] for each one.

Collectives™ on Stack Overflow

Select elements of a numpy array based on the elements of a second array

2 Answers 2

1 Comment

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related