如何高效按行索引从numpy数组或pandas DataFrame取值？

Great question! The problem with using np.take directly is that it operates on the flattened version of your array, not applying the index to each column individually. We can fix this with numpy's broadcasting and reshaping capabilities to create a fully vectorized solution—no loops needed, and it’ll be way faster for larger datasets.

Method 1: np.take with Dimension Expansion #

This is the most concise approach. We’ll add an extra dimension to the original array so take can apply the index to each column separately, then reshape the result to your target shape:

import numpy as np

arr = np.asarray([[0, 1, 2, 4], [1, 2, 3, 4], [2, 3, 4, 5]])
indices = np.asarray([[-1, 0, 1], [-1, -1, 0]]).T

# Add a third dimension to arr (shape becomes (3,4,1))
arr_expanded = arr[:, :, np.newaxis]
# Apply indices to the row axis (axis=0) for each column
temp_result = arr_expanded.take(indices, axis=0)
# Reshape to (3, 4*2) by flattening the last two dimensions
res_efficient = temp_result.reshape(arr.shape[0], -1)

print(f"efficient result:\n{res_efficient}")

Output #

efficient result:
[[2 2 3 3 4 4 5 5]
 [0 2 1 3 2 4 4 5]
 [1 0 2 1 3 2 4 4]]

How It Works #

• arr[:, :, np.newaxis] turns your (3,4) array into (3,4,1), treating each column as a separate "slice" along the new third dimension.
• take(indices, axis=0) applies the index array to each of these slices, resulting in a (3,4,2) array—each original column now has 2 columns from the index lookup.
• reshape(arr.shape[0], -1) flattens the last two dimensions into one, giving you the desired (3,8) shape.

Method 2: Advanced Indexing #

If you prefer using numpy’s advanced indexing syntax, you can explicitly construct row and column index arrays to target the exact elements you need:

# Create row indices (shape: (3,1,1))
rows = np.arange(arr.shape[0])[:, None, None]
# Create column indices (shape: (1,4,1))
cols = np.arange(arr.shape[1])[None, :, None]
# Use broadcasting to match dimensions and select elements
res_advanced = arr[rows, cols, indices[:, None, :]].reshape(arr.shape[0], -1)

This achieves the same result as Method 1, using explicit index matching instead of take.

Performance Benefits #

Vectorized operations like these are drastically faster than looping, especially as your dataset grows. For example, if you have an array with 1000 rows and 1000 columns, the loop approach would run 1000 separate take calls, while the vectorized method does it in one step—you’ll see a difference of seconds vs. milliseconds.

For Pandas DataFrames #

If you’re working with pd.DataFrame instead of numpy arrays, just convert the DataFrame to a numpy array, run the vectorized operation, then convert back:

import pandas as pd

df = pd.DataFrame(arr)
indices_df = pd.DataFrame(indices)

# Convert to numpy, process, then convert back
result_np = df.values[:, :, np.newaxis].take(indices_df.values, axis=0).reshape(df.shape[0], -1)
result_df = pd.DataFrame(result_np)

This keeps the convenience of pandas while leveraging numpy’s speed.

内容的提问来源于stack exchange，提问作者DuFei