Python numpy 2D array indexing

Solution 1:

In [1]: import numpy as np
In [2]: a = np.array([[2,0],[3,0],[3,1],[5,0],[5,1],[5,2]])
In [3]: b = np.zeros((6,3), dtype='int32')

In [4]: b[a[:,0], a[:,1]] = 10

In [5]: b
Out[5]: 
array([[ 0,  0,  0],
       [ 0,  0,  0],
       [10,  0,  0],
       [10, 10,  0],
       [ 0,  0,  0],
       [10, 10, 10]])

Why it works:

If you index b with two numpy arrays in an assignment,

b[x, y] = z

then think of NumPy as moving simultaneously over each element of x and each element of y and each element of z (let's call them xval, yval and zval), and assigning to b[xval, yval] the value zval. When z is a constant, "moving over z just returns the same value each time.

That's what we want, with x being the first column of a and y being the second column of a. Thus, choose x = a[:, 0], and y = a[:, 1].

b[a[:,0], a[:,1]] = 10

Why b[a] = 10 does not work

When you write b[a], think of NumPy as creating a new array by moving over each element of a, (let's call each one idx) and placing in the new array the value of b[idx] at the location of idx in a.

idx is a value in a. So it is an int32. b is of shape (6,3), so b[idx] is a row of b of shape (3,). For example, when idx is

In [37]: a[1,1]
Out[37]: 0

b[a[1,1]] is

In [38]: b[a[1,1]]
Out[38]: array([0, 0, 0])

So

In [33]: b[a].shape
Out[33]: (6, 2, 3)

So let's repeat: NumPy is creating a new array by moving over each element of a and placing in the new array the value of b[idx] at the location of idx in a. As idx moves over a, an array of shape (6,2) would be created. But since b[idx] is itself of shape (3,), at each location in the (6,2)-shaped array, a (3,)-shaped value is being placed. The result is an array of shape (6,2,3).

Now, when you make an assignment like

b[a] = 10

a temporary array of shape (6,2,3) with values b[a] is created, then the assignment is performed. Since 10 is a constant, this assignment places the value 10 at each location in the (6,2,3)-shaped array. Then the values from the temporary array are reassigned back to b. See reference to docs. Thus the values in the (6,2,3)-shaped array are copied back to the (6,3)-shaped b array. Values overwrite each other. But the main point is you do not obtain the assignments you desire.