如何将数组中除argmax对应元素外全部置零？求更优实现方案

Great question! Your list comprehension works for small 1D lists, but when dealing with larger arrays or multi-dimensional data, vectorized NumPy operations are far more efficient—they leverage optimized C under the hood instead of slow Python-level loops. Let’s break down the best approaches for both 1D and multi-dimensional scenarios:

1D Arrays: Simple Vectorized Assignment #

For a 1D NumPy array, replace your list comprehension with this much faster method:

import numpy as np

a = np.array([1, 2, 3, 4, 5])
b = np.zeros_like(a)  # Create an array of zeros matching a's shape and dtype
max_idx = np.argmax(a)
b[max_idx] = a[max_idx]

This produces array([0, 0, 0, 0, 5]) just like your original code, but runs orders of magnitude faster for large arrays—test it with %timeit to see the stark difference!

Handling Multiple Max Values #

If your array has multiple instances of the maximum value and you want to keep all of them (not just the first occurrence), use boolean indexing instead:

a = np.array([5, 2, 5, 4])
b = np.zeros_like(a)
b[a == a.max()] = a[a == a.max()]

Result: array([5, 0, 5, 0])—still fully vectorized and efficient.

Multi-Dimensional Arrays #

For 2D, 3D, or higher-dimensional arrays, we need to handle index conversion from flattened to multi-dimensional coordinates:

Single Global Maximum #

To keep only the single global maximum (the first occurrence if there are ties):

# 2D example array
a = np.array([[1, 5], [3, 2]])
b = np.zeros_like(a)

# Get flattened index of the maximum, then convert to multi-dimensional indices
flat_max_idx = np.argmax(a)
multi_max_idx = np.unravel_index(flat_max_idx, a.shape)

b[multi_max_idx] = a[multi_max_idx]

Result: array([[0, 5], [0, 0]])

Maxima Along a Specific Axis #

If you want to keep the maximum value for each row (or column) instead of a global maximum, use advanced indexing:

a = np.array([[1, 5], [3, 2]])
b = np.zeros_like(a)

# Get max indices for each row (axis=1)
row_max_indices = np.argmax(a, axis=1)

# Use advanced indexing to set the max values in each row
b[np.arange(a.shape[0]), row_max_indices] = a[np.arange(a.shape[0]), row_max_indices]

Result: array([[0, 5], [3, 0]])

Why This Is Better Than Your Original Code #

• Speed: List comprehensions iterate over each element in Python, which is slow for large arrays. NumPy’s vectorized operations run in optimized C code, so you’ll see massive speed gains (e.g., 100x faster for arrays with 1M elements).
• Scalability: These methods work seamlessly for any number of dimensions, without rewriting core logic.
• Readability: The intent is clearer—you’re explicitly creating a zero array and setting only the maximum value(s), rather than looping with a conditional.

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