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
JavaScript
x
14
14
1
A = [[0,1],[3,3]] -> I would want B to be
2
3
[[[1,0],
4
[0,0]],
5
6
[0,1],
7
[0,0]],
8
9
[0,0],
10
[0,1]],
11
12
[0,0],
13
[0,2]]]
14
How do you do this in numpy efficiently?
Advertisement
Answer
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:
JavaScript
1
2
1
B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment
2
For example:
JavaScript
1
19
19
1
A = np.array([[0,1],[3,3]])
2
B = np.zeros((4,2,2), dtype=int)
3
increment = 1
4
5
B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment
6
7
>>> B
8
array([[[1, 0],
9
[0, 0]],
10
11
[[0, 1],
12
[0, 0]],
13
14
[[0, 0],
15
[0, 0]],
16
17
[[0, 0],
18
[1, 1]]])
19
Another way of doing the same thing is:
JavaScript
1
6
1
w, h = A.shape
2
indices = (A.ravel(),) + tuple(np.mgrid[:w, :h].reshape(2, -1))
3
4
# then
5
B[indices] += increment
6