4

How can I create a NumPy array B which is a sub-array of a NumPy array A, by specifying which rows and columns (indicated by x and y respectively) are to be included?

For example:

A = numpy.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
x = [0, 2]
y = [1, 3, 4]
B = # Do something .....

Should give the output:

>>> B
array([[2, 4, 5], [12, 14, 15]])
Alex Riley
  • 169,130
  • 45
  • 262
  • 238
Karnivaurus
  • 22,823
  • 57
  • 147
  • 247

3 Answers3

6

The best way to do this is to use the ix_ function: see the answer by MSeifert for details.

Alternatively, you could use chain the indexing operations using x and y:

>>> A[x][:,y]
array([[ 2,  4,  5],
       [12, 14, 15]])

First x is used to select the rows of A. Next, [:,y] picks out the columns of the subarray specified by the elements of y.

The chaining is symmetric in this case: you can also choose the columns first with A[:,y][x] if you wish.

Alex Riley
  • 169,130
  • 45
  • 262
  • 238
  • Note that this is not a new array, but a view of the original one, using the same data in memory. – sebix Oct 22 '14 at 17:21
  • This is actually a new array - fancy indexing (i.e using an array of values, rather than a single value or slice) triggers copying. – Alex Riley Feb 10 '18 at 09:36
2

You can use np.ix_ which allows to broadcast the integer index arrays:

>>> A = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
>>> x = [0, 2]
>>> y = [1, 3, 4]
>>> A[np.ix_(x, y)]
array([[ 2,  4,  5],
       [12, 14, 15]])

From the documentation the ix_ function was designed so

[...] one can quickly construct index arrays that will index the cross product. a[np.ix_([1,3],[2,5])] returns the array [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

MSeifert
  • 145,886
  • 38
  • 333
  • 352
0

Here's a super verbose way to get what you want:

import numpy as np
a = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
x = [0, 2]
y = [1,3,4]

a2 = a.tolist()
a3 = [[l for k,l in enumerate(j) if k in y] for i,j in enumerate(a2) if i in x]
b = np.array(a3)

But please follow @ajcr answer:

import numpy as np
a = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
x = [0, 2]
y = [1,3,4]
a[x][:,y]
alvas
  • 115,346
  • 109
  • 446
  • 738