CODEWITHSUNDEEP
pythonnumpy
generic_user
generic_user

Reputation: 3562

indexing rows and columns in numpy

a = np.array(list(range(16).reshape((4,4))
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

Say I want the middle square. It'd seem reasonable to do this:

a[[1,2],[1,2]]

but I get this:

array([5, 10])

This works, but seems inelegant:

a[[1,2],:][:,[1,2]]
array([[5, 6],
       [9, 10]])

So my questions are:

  1. Why is it this way? What premises are required to make the implemented way sensible?
  2. Is there a canonical way to select along more than one index at once?

Upvotes: 4

Views: 1762

Answers (3)

Kevin
Kevin

Reputation: 3358

You could use multiple np.take to select indices from multiple axes

a = np.arange(16).reshape((4, 4))
idx = np.array([1,2])
np.take(np.take(a, idx, axis=1), idx, axis=0)

Or (slightly more readable)

a.take(idx, axis=1).take(idx, axis=0)

Output:

array([[ 5,  6],
       [ 9, 10]])

np.take also allows you to conveniently wrap around out-of-bound indices and such.

Upvotes: 0

Stefan B
Stefan B

Reputation: 1677

You can do both slicing operations at once instead of creating a view and indexing that again:

import numpy as np

a = np.arange(16).reshape((4, 4))

# preferred if possible
print(a[1:3, 1:3])
# [[ 5  6]
#  [ 9 10]]

# otherwise add a second dimension to the first index to make it broadcastable
index1 = np.asarray([1, 2])
index2 = np.asarray([1, 2])
print(a[index1[:, None], index2])
# [[ 5  6]
#  [ 9 10]]

Upvotes: 1

Quang Hoang
Quang Hoang

Reputation: 150745

I think you can read more details on advanced indexing. Basically, when you slice the array by lists/arrays, the arrays will be broadcast and iterate together.

In your case, you can do:

idx = np.array([1,3])
a[idx,idx[:,None]]

Or as in the doc above:

a[np.ix_(idx, idx)]

Output:

array([[ 5, 13],
       [ 7, 15]])

Upvotes: 1

Related Questions