How to Extract Variable Column Slices Per Row in a 2D NumPy Array?

Below are a couple of ways to get different slices on each row.

1. Using Advanced (Fancy) Indexing

You can compute row and column index arrays. For example, to get for row 0 the slice starting at idx[0] and for row 1 the slice starting at idx[1], you can do:

PYTHONCopied
import numpy as np

x = np.arange(12).reshape(3, 4)
idx = np.array([0, 1])  # starting indices for each row

# Get the row indices (for the first two rows)
rows = np.arange(2)[:, None]

# Create a corresponding 2D array for the columns
cols = idx[:, None] + np.arange(2)

result = x[:2][rows, cols]
print(result)

This outputs:

[[0 1]
 [5 6]]

Here, we first build a 2D array cols where each row contains the two consecutive column indices we want to extract. Then we use advanced indexing with a 2D coordinates array.

2. Using a List Comprehension

If you prefer simpler code or have a variable slice length in each row, you can loop over the rows in Python:

PYTHONCopied
import numpy as np

x = np.arange(12).reshape(3, 4)
idx = np.array([0, 1])  # starting indices for each row

result = np.array([row[i:i+2] for row, i in zip(x[:2], idx)])
print(result)

The output will again be:

[[0 1]
 [5 6]]

Performance Considerations

• Advanced indexing: Implemented in C, so it is typically faster and more efficient, especially for large arrays. However, note that advanced indexing always returns a copy of the data.
• List comprehension: Although more intuitive, it uses a Python loop, which may be slower for large arrays.

Your choice depends on your specific needs and data size. For high-performance applications, the advanced indexing method is usually preferable.