如何替换n维NumPy数组中指定轴索引位置的元素

If you need to replace all elements at a specific index along a given axis in an n-dimensional NumPy array, NumPy's flexible indexing system makes this straightforward. Here's how to do it, using your 3D array example as a reference:

Step 1: Understand the Indexing Approach #

We can construct an index tuple that targets the exact position we want:

• For every axis except the one we're targeting, use : (or slice(None)) to select all elements.
• For the target axis, specify the index we want to replace.

Step 2: Example Implementation #

Let's walk through your sample scenario:

First, create the initial 3D array of ones:

import numpy as np

a = np.ones((2, 2, 2))
print("Original array:")
print(a)

Output:

Original array:
array([[[1., 1.],
        [1., 1.]],

       [[1., 1.],
        [1., 1.]]])

Now, replace all elements at axis=1, index=0 with 0:

# Build the index tuple: target axis 1 at index 0, others select all
indices = tuple([slice(None) if i != 1 else 0 for i in range(a.ndim)])
a[indices] = 0

print("\nModified array:")
print(a)

Output:

Modified array:
array([[[0., 0.],
        [1., 1.]],

       [[0., 0.],
        [1., 1.]]])

Step 3: Reusable Function #

To make this easier to use across different arrays, wrap the logic in a function:

def replace_axis_elements(arr, axis, index, value):
    # Create index tuple targeting the specified axis and index
    indices = tuple([slice(None) if i != axis else index for i in range(arr.ndim)])
    arr[indices] = value
    return arr

You can use it like this:

a = np.ones((2, 2, 2))
modified_a = replace_axis_elements(a, axis=1, index=0, value=0)
print(modified_a)

This approach works for any n-dimensional array—just adjust the axis, index, and value parameters as needed.

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