How to select one element in each row of a NumPy array by column indices?

Given a NumPy array, we have to select one element in each row from it by column indices.
By Pranit Sharma Last updated : December 26, 2023

Problem statement

Suppose that we are given a numpy array and an array of indices to be selected and we need to select one element in each row of this array using the array of indices.

NumPy Array | Selecting one element in each row by column indices

For this purpose, we can choose from the given array using numpy.choose() which constructs an array from an index array and a set of arrays to choose from. However, we may first need to transpose the input array to match the dimensions.

Let us understand with the help of an example,

Python code to select one element in each row of a NumPy array by column indices

# Import numpy
import numpy as np

# Creating a numpy array
arr = np.array([[ 3, 14],[12,  5],[75, 50]])

# Display Original array
print("Original array:\n",arr,"\n\n")

# Array of indices
id_arr = [0, 1, 1]

# Using np.choose to select elements of input array
res = np.choose(id_arr, arr.T)

# Display result
print("Result:\n",res)