CODEWITHSUNDEEP
pythonarrayslistnumpy
Tsuyoshi
Tsuyoshi

Reputation: 27

Extract specific array from multi dimensional array with index sets

Suppose inputs are like below.

indexSet1 = [0,1,2]
indexSet2 = [1,2]
A = [[1,2,3],[4,5,6],[7,8,9]]

Then I want to get a matrix whose height is 3 and width is 2 respectively and elements corresponds to indexSet1's value and indexSet2's one. In short, I want to a array [[2,3],[4,5],[7,8]] from A by indexSet1 and indexSet2.

When I code like below, I can not get my desire result.

>>> import numpy as np
>>> np.array(A)[np.array(indexSet1),np.array(indexSet2)]
array([5, 9])

Can anyone know wise methods? Sorry for my poor English. And thank you in advance.

Upvotes: 0

Views: 68

Answers (2)

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 251146

In NumPy you can do something like this:

>>> A = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> A[np.array(indexSet1)[:, None], indexSet2]
array([[2, 3],
       [5, 6],
       [8, 9]])

Upvotes: 1

falsetru
falsetru

Reputation: 369424

Using nested list comprehension:

>>> indexSet1 = [0,1,2]
>>> indexSet2 = [1,2]
>>> A = [[1,2,3],[4,5,6],[7,8,9]]
>>> [[A[i][j] for j in indexSet2] for i in indexSet1]
[[2, 3], [5, 6], [8, 9]]

Upvotes: 1

Related Questions