CODEWITHSUNDEEP
pythonarraysnumpy
fgregg
fgregg

Reputation: 3249

Fancy indexing with assignment for numpy array

I have a 2 dimensional numpy array:

A = np.zeros(16).reshape(4,4)

I want the (1, 1), (1,3), (3,1), and (3,3) cells to take a value of 1.

A[[1,3], [1:3]] = 1

Only assigns 1 to (1,1) and (3,3). 

A[[1,3], :][:, [1, 3]] = 1

Does not work because it make a copy of the data not a view. What's the right way to do this?

Upvotes: 7

Views: 2282

Answers (3)

Peter
Peter

Reputation: 13485

>>> A = np.zeros(16).reshape(4,4)
>>> A[[[1],[3]], [[1, 3]]] = 1
>>> A
array([[ 0.,  0.,  0.,  0.],
   [ 0.,  1.,  0.,  1.],
   [ 0.,  0.,  0.,  0.],
   [ 0.,  1.,  0.,  1.]])

What's happening here is that the row indices [[1],[3]] have shape (2, 1), so they can be broadcast against the column indices [[1, 3]], which have shape (1, 2), resulting in shapes (2, 2) arrays of row/column indices. So we can think of the above as boing a short form for:

>>> A[[[1, 1],[3, 3]], [[1, 3], [1, 3]]] = 1

Which produces the same result.

Upvotes: 1

senderle
senderle

Reputation: 150977

Normally, when you use arrays to index another array this way, numpy expects each array R, C, and so on, to have the same shape. So for example, suppose you want to extract the nonzero values from this array:

>>> a
array([[1, 3, 0, 0],
       [0, 0, 0, 0],
       [2, 4, 0, 0],
       [0, 0, 0, 0]])

You would create an array of row indices R:

>>> R
array([[0, 2],
       [0, 2]])

And an array of column indices C:

>>> C
array([[0, 0],
       [1, 1]])

And pass them in like so:

>>> a[R, C]
array([[1, 2],
       [3, 4]])

Note that these can be in whatever shape you like -- the output will take the same shape:

>>> RR
array([0, 2, 0, 2])
>>> CC
array([0, 0, 1, 1])
>>> a[RR, CC]
array([1, 2, 3, 4])

However, if your index arrays have repetitions in them, then you can save yourself some trouble by using broadcasting. The only catch is that the resulting arrays must be broadcastable. That means you need to explicitly give them extra dimensions. I'll use slice syntax to preserve the extra dimension.

>>> r = R[0:1,:]
>>> c = C[:,0:1]
>>> r
array([[0, 2]])
>>> c
array([[0],
       [1]])
>>> a[r, c]
array([[1, 2],
       [3, 4]])

If you don't explicitly give them extra dimensions, numpy does its best to make sense of what you've passed, but it won't always work as expected. Replacing the 0:1 slices with 0 removes the extra dimension:

>>> rr = r[0,:]
>>> cc = c[:,0]
>>> rr
array([0, 2])
>>> cc
array([0, 1])
>>> a[rr, c]
array([[1, 2],
       [3, 4]])
>>> a[r, cc]
array([[1, 4]])
>>> a[rr, cc]
array([1, 4])

The first one (a[rr, c]) works because numpy can tell from the shape of c that it should broadcast. But the other two are ambiguous, so numpy assumes that you don't mean them to be broadcasted.

However, note that numpy does also provide a shortcut for making flat arrays broadcastable. This can come in handy!

>>> a[numpy.ix_(rr, cc)]
array([[1, 3],
       [2, 4]])

Upvotes: 4

shx2
shx2

Reputation: 64318

Use slices with step=2:

A[1::2,1::2] = 1

Or explictly pass all the indexes:

A[[1,1,3,3],[1,3,1,3]] = 1

Upvotes: 5

Related Questions