How to use the value of a 2d array as an index to a 3d array in numpy?

Question

I have a 2d array, A, with shape (n x m), where each element of the array at position (i,j) holds a third value k. I want to increment a 3d array with dimensions nxmxl at position (k,i,j) based on the 2d array value and position.

So for example if

A = [[0,1],[3,3]] -> I would want B to be 

[[[1,0],
  [0,0]],

  [0,1],
  [0,0]],

  [0,0],
  [0,1]],

  [0,0],
  [0,2]]]

How do you do this in numpy efficiently?

Shows how you'd do this in-efficiently. The action isn't obvious from your example and word description. — hpaulj
– hpaulj, Commented Apr 9, 2022 at 15:00
Your B has shape (4,2,2). A values could be indices of the first dimension (4), but not either of the others. Your problem might be easier if A was 1d, and B 2d. The mix of 2 and 3d probably isn't essential. — hpaulj
– hpaulj, Commented Apr 9, 2022 at 15:36
@hpaulj, sorry, the first axis is implied from max value of the array, in this case 3. — Mingssdjk
– Mingssdjk, Commented Apr 9, 2022 at 17:33

hpaulj · Accepted Answer · 2022-04-09 16:14:57Z

1

I can produce your B with:

In [208]: res = np.zeros((4,2,2),int)
In [209]: res.reshape(4,4)[np.arange(4), A.ravel()] = [1,1,1,2]
In [210]: res
Out[210]: 
array([[[1, 0],
        [0, 0]],

       [[0, 1],
        [0, 0]],

       [[0, 0],
        [0, 1]],

       [[0, 0],
        [0, 2]]])

I use the reshape because A values look more like indices of

In [211]: res.reshape(4,4)
Out[211]: 
array([[1, 0, 0, 0],
       [0, 1, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 2]])

answered Apr 9, 2022 at 16:14

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

1 Comment

Pierre D Over a year ago

This is so cool. You got a different read of the question than I did (and yours might be the correct one), but I didn't know you could use .reshape() as an l-value. After years of usage, I keep marveling at how elegant and powerful numpy is.

Pierre D · Accepted Answer · 2022-04-09 16:15:07Z

0

The question is somewhat ambiguous, but if the intent is to increment some unknown array B at indices (0,0,0), (1,0,1), (3,1,0), and (3,1,1), then the following should be fine:

B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment

For example:

A = np.array([[0,1],[3,3]])
B = np.zeros((4,2,2), dtype=int)
increment = 1

B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment

>>> B
array([[[1, 0],
        [0, 0]],

       [[0, 1],
        [0, 0]],

       [[0, 0],
        [0, 0]],

       [[0, 0],
        [1, 1]]])

Another way of doing the same thing is:

w, h = A.shape
indices = (A.ravel(),) + tuple(np.mgrid[:w, :h].reshape(2, -1))

# then
B[indices] += increment

edited Apr 9, 2022 at 16:15

answered Apr 9, 2022 at 16:04

Pierre D

26.6k8 gold badges71 silver badges108 bronze badges

Collectives™ on Stack Overflow

How to use the value of a 2d array as an index to a 3d array in numpy?

2 Answers 2

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related