Currently, np.insert returns the modified array, but the caller loses the exact indices where the new elements were placed. In complex data pipelines, machine learning preprocessing, or systems requiring state tracking (like undo/redo buffers), knowing exactly where data landed is critical. Calculating this manually after the fact is error-prone, especially with multi-dimensional arrays or variable-length insertions.
I propose adding a return_indices=False keyword argument to numpy.insert. When set to True, the function returns a tuple: (inserted_array, indices_array).
Example Usage
import numpy as np
arr = np.array([10, 20, 30, 40, 50])
# Insert multiple values, tracking their final location
new_arr, indices = np.insert(arr, 2, [99, 99, 99], return_indices=True)
print(new_arr)
# [10, 20, 99, 99, 99, 30, 40, 50]
print(indices)
# [2, 3, 4]
Implementation Thoughts & Multi-Dimensional Handling For 1D arrays, the logic is straightforward (calculating the offset based on np.size(values)).
For N-dimensional arrays where an axis is specified, return_indices should return the indices along that specific axis where the new sub-arrays were placed. If axis=None (where the array is flattened before insertion), it would return the 1D indices of the flattened result.
Writing custom wrapper functions to handle this in user-land is possible, but it creates fragmented, unoptimized boilerplate across the ecosystem for what is fundamentally a core array operation. I’d love to hear the core team’s thoughts on whether this fits into the main API, and I am happy to work on a PR if there is interest.