[FIXED] How to index 2d Numpy array with 2d integer array

Issue

I am looking to index a 2d array with a 2d array of integers. The goal is to index each column separately, however, i can't seem to find the right syntax for this.

I understand why the naive x[s] does not work, but can only get it to work in an inefficient list comprehension.

import numpy as np
x = np.random.normal(size=(10, 5))
s = np.argsort(x, 0)

What i hoped would work, but shape is (10, 5, 5)

x[s] 

Desired outcome, shape: (10, 5)

np.vstack([x[s[:, i], i] for i in range(len(x[0]))]).T # -> shape is (10, 5)

Solution

You can use s to index into rows of x if you also specify the columns:

# The list comprehension that you know works correctly
correct_result = np.vstack([x[s[:, i], i] for i in range(len(x[0]))]).T

# Indexing by also specifying column indices
result = x[s, range(x.shape[1])]

print((result == correct_result).all()) # True

Answered By - pho