How can I extract row-specific slices from a 2D NumPy array?

Below are three common options to achieve “different slices per row” when the starting columns vary by row. Note that basic slicing (as in your code) cannot accept an array for the slice boundaries. Each option has its own tradeoffs:

---

Option 1: List Comprehension

You can loop over each row and extract the corresponding slice. This is clear and easy to read but runs in Python (i.e. a Python-level loop):

PYTHONCopied
import numpy as np

x = np.arange(12).reshape(3, 4)
idx = np.array([0, 1])
# For each row in the first two rows, slice the columns starting at idx[i]
y = np.array([x[i, idx[i]: idx[i] + 2] for i in range(2)])
print(y)
# Output:
# [[0 1]
#  [5 6]]

---

Option 2: Advanced (Fancy) Indexing via Broadcasting

You can build index arrays that pick the desired elements without an explicit Python loop. This method is fully vectorized but note that advanced indexing always returns a copy:

PYTHONCopied
import numpy as np

x = np.arange(12).reshape(3, 4)
idx = np.array([0, 1])
rows = np.arange(len(idx))           # rows: [0, 1]
# For each row create columns that go from idx[i] to idx[i]+2
cols = idx[:, None] + np.arange(2)     # cols becomes:
# [[0, 1],
#  [1, 2]]
y = x[rows[:, None], cols]             # advanced indexing with broadcasting
print(y)
# Output:
# [[0 1]
#  [5 6]]

---

Option 3: Using np.take_along_axis

With np.take_along_axis, you can create an index array corresponding to the slice for each row and then extract the values. This is also vectorized like advanced indexing (and produces a copy):

PYTHONCopied
import numpy as np

x = np.arange(12).reshape(3, 4)
idx = np.array([0, 1])
rows = np.arange(len(idx))[:, None]  # shape (2, 1)
# Create an index for two consecutive columns per row
cols = idx[:, None] + np.arange(2)     # shape (2, 2)
# np.take_along_axis requires that the index array has the same shape as the desired output
y = np.take_along_axis(x[:2], cols, axis=1)
print(y)
# Output:
# [[0 1]
#  [5 6]]

---

Performance Implications

• List Comprehension:
  May be slower for very large arrays because it involves a Python-level loop and creates many small arrays before combining them.

• Advanced Indexing / np.take_along_axis:
  These are vectorized operations that are implemented in C, so they generally perform better on large arrays. However, they create copies of data (no view is returned), which might have memory implications if the slices are very large.

---

Each of these methods has its use depending on your performance needs and code clarity requirements.