numpy多维数组索引

Question

I have a strange issue when I am doing indexing in a numpy multidimensional array. So, I have an array of shape (4, 882). I have another array called matches of shape (276, 2). This match array holds indexes of valid entries in my original multidimensional array. What I would like to do is select the first 2 rows and all the columns matching the indexes in the match array. So, I do something as follows:

import numpy as np
k = get_array()  # This has shape (4, 882)
m = get_match()  # This has shape (276, 2)

s = k[[1, 0], m[:, 0]]

This raises the error:

ValueError: shape mismatch: objects cannot be broadcast to a single shape

However, when I do:

s = k[[1, 0], :][:, m[:, 0]]

This works fine. So, this is in effect selecting the subset of rows first and then the columns but I am not sure why my first attempt is wrong. Also, doing:

s = k[[1, 0], :]

of course works.

Answer 1

The error message is a little confusing because the shape mismatch isn't between k and m . It's between [1, 0] and m[:,0] . Here are three ways to fix it, using the following (more easily visualizable) arrays:

>>> k
array([[ 0,  1,  2,  3,  4,  5,  6,  7],
       [ 8,  9, 10, 11, 12, 13, 14, 15],
       [16, 17, 18, 19, 20, 21, 22, 23]])
>>> m
array([[0, 1],
       [2, 3],
       [4, 5],
       [6, 7]])

1. Change the indexing list to a 2-d list with a shape that can be broadcasted against m :

    >>> k[[[1], [0]],m[:, 0]] array([[ 8, 10, 12, 14], [ 0, 2, 4, 6]]) 

2. Convert it to an array and reshape the array:

    >>> k[numpy.array([1, 0]).reshape(2, 1),m[:, 0]] array([[ 8, 10, 12, 14], [ 0, 2, 4, 6]]) 

3. Convert it to an array and slice it using numpy.newaxis , aka None :

    >>> k[numpy.array([1, 0])[:,None],m[:, 0]] array([[ 8, 10, 12, 14], [ 0, 2, 4, 6]]) 

There are many others, but these are the three that come most readily to mind.

You're getting an error because numpy needs indices passed in this way to have the same shape, or to be broadcastable to the same shape. By making the [1, 0] list a "column array," you make them broadcastable. The same thing happens when you try to multiply m by [1, 0] :

>>> m[:,0] * [0, 1]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: shape mismatch: objects cannot be broadcast to a single shape

And all the same fixes apply. For example:

>>> m[:,0] * [[0], [1]]
array([[0, 0, 0, 0],
       [0, 2, 4, 6]])

Finally, note that you could also fix it by passing a differently-shaped slice of m -- observe that the output is transposed:

>>> k[[1, 0],m[:, 0:1]]
array([[ 8,  0],
       [10,  2],
       [12,  4],
       [14,  6]])