0
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
index = np.array([0,1,2])
b = a[index]
b

I expect result like: [1,5,9]

I found some solutions:

a = np.array([[1,2,3],[4,5,6],[7,8,9]])
index = np.array([1,0,2])
b = a[:,index]
np.diagonal(b)

a = np.array([[1,2,3],[4,5,6],[7,8,9]])
index = np.array([1,0,2])
b = a[np.arange(3),index]
b

Probably

Dmitry Sokolov
  • 1,303
  • 1
  • 17
  • 29
  • If the index is always a permutation of `np.arange(a.shape[1])`, you can use the suggested approach of `np.diag(a[index.argsort()])`. – Mad Physicist Dec 24 '21 at 11:50

1 Answers1

2

You need to work with the array indexes to get your desired output.

import numpy as np

your_array = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
idx, jdx = (0, 1, 2), (0, 1, 2)
sub_array = your_array[idx, jdx]

Out[18]: array([1, 5, 9])
Hommes
  • 286
  • 1
  • 10