CODEWITHSUNDEEP
pythonarraysnumpy
Karnivaurus
Karnivaurus

Reputation: 24121

Create NumPy array from another array by specifying rows and columns

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]])

Upvotes: 4

Views: 19089

Answers (3)

Alex Riley
Alex Riley

Reputation: 176830

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.

Upvotes: 6

MSeifert
MSeifert

Reputation: 152657

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]]].

Upvotes: 2

alvas
alvas

Reputation: 122072

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]

Upvotes: 0

Related Questions