为何np.argwhere输出形状与输入不匹配？1D数组变2D的原因及优势

Let's break this down clearly using your example first. Here's a quick recap of what you observed:

import numpy as np

# 1D input array
arr = np.arange(0,20)
print(arr.shape)  # Output: (20,)

# Apply argwhere with a condition
result = np.argwhere(arr < 10)
print(result)
# Output: [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9]]
print(result.shape)  # Output: (10, 1)

Why the 2D output? #

The key driver here is consistency across all input dimensions. np.argwhere isn't built just for 1D arrays—it's designed to handle arrays of any number of dimensions, and it uses a uniform output format regardless of the input's shape.

Take a 2D array example to see the pattern:

arr_2d = np.arange(0,20).reshape(4,5)
result_2d = np.argwhere(arr_2d < 10)
print(result_2d.shape)  # Output: (10, 2)

Here, each row in result_2d gives the (row, column) indices of elements that match the condition. For a 1D array, it follows the exact same logic: each row holds the single index (from the 1D axis) of a matching element. This way, no matter if your input is 1D, 2D, or even 3D+, you always get an N x D array where N is the number of matches and D is the number of dimensions in the input.

What's the benefit of this design? #

This choice makes your code more robust and intuitive, with three main advantages:

• No special-case handling: You don't have to write separate logic to process results from 1D vs higher-dimensional inputs. The output structure is predictable, so you can reuse the same code across different array shapes without tweaks.
• Seamless indexing: The 2D format plays perfectly with NumPy's advanced indexing. For your 1D example, you can easily extract matching elements without reshaping the result:

  matching_elements = arr[result[:, 0]]
  print(matching_elements)  # Output: [0 1 2 3 4 5 6 7 8 9]
  

  For 2D inputs, you can use arr_2d[tuple(result_2d.T)] to get all matches—same logical flow, no extra work.
• Clear index context: Each row explicitly represents a complete index for the original array. For 1D, it's [index]; for 2D, [row, col]; for 3D, [depth, row, col]. This removes ambiguity about where each matching element sits in the array's structure.

If argwhere returned a 1D array for 1D inputs, you'd constantly have to check the input's dimensions to handle the result correctly—adding unnecessary complexity to your code.

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