如何筛选Numpy数组中邻域元素大于指定值的元素索引？

To solve this problem, we need to identify elements in a NumPy array where at least one adjacent neighbor (in any dimension) is greater than or equal to a given value x. The solution uses vectorized operations for efficiency, making it suitable for large arrays and compatible with arbitrary dimensions.

Approach #

The core idea is to check each element's neighbors in every direction (positive and negative) along each axis. We do this by:

1. For each axis, creating shifted versions of the array that represent the value of neighbors in the positive (e.g., right, down) and negative (e.g., left, up) directions.
2. Converting these shifted arrays into boolean masks indicating if the neighbor is >= x.
3. Combining all these masks using logical OR to get the final result (True if any neighbor meets the condition).

Solution Code #

import numpy as np

def neighboringIndicesLargerThan(arr, x):
    result = np.zeros_like(arr, dtype=bool)
    ndim = arr.ndim
    
    for axis in range(ndim):
        # Check neighbor in the negative direction (i-1) along the current axis
        slice_neg = tuple(slice(None, -1) if i == axis else slice(None) for i in range(ndim))
        pad_neg = tuple((1, 0) if i == axis else (0, 0) for i in range(ndim))
        neg_neighbor_mask = np.pad(arr[slice_neg] >= x, pad_neg, mode='constant', constant_values=False)
        result = np.logical_or(result, neg_neighbor_mask)
        
        # Check neighbor in the positive direction (i+1) along the current axis
        slice_pos = tuple(slice(1, None) if i == axis else slice(None) for i in range(ndim))
        pad_pos = tuple((0, 1) if i == axis else (0, 0) for i in range(ndim))
        pos_neighbor_mask = np.pad(arr[slice_pos] >= x, pad_pos, mode='constant', constant_values=False)
        result = np.logical_or(result, pos_neighbor_mask)
    
    return result

Explanation #

• Initialization: We start with a boolean array of False values with the same shape as the input array.
• Loop through each axis: For each dimension in the array:
  • Negative direction: We take all elements except the last one along the axis, then pad a slice of False at the start to align with the original array shape. This mask represents whether the neighbor in the negative direction (e.g., above, left) is >= x.
  • Positive direction: We take all elements except the first one along the axis, then pad a slice of False at the end. This mask represents whether the neighbor in the positive direction (e.g., below, right) is >= x.
• Combine masks: We use np.logical_or to combine each neighbor mask into the result array, so any element with at least one qualifying neighbor becomes True.

Example Usage #

2D Array Test

test = np.arange(4**2).reshape((4,4))
print("Input array:")
print(test)
print("\nResult (neighbor >=9):")
print(neighboringIndicesLargerThan(test, 9))

Output:

Input array:
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

Result (neighbor >=9):
[[False False False False]
 [False  True  True  True]
 [ True  True  True  True]
 [ True  True  True  True]]

3D Array Test

test3d = np.arange(8).reshape(2,2,2)
print("3D Input array:")
print(test3d)
print("\nResult (neighbor >=5):")
print(neighboringIndicesLargerThan(test3d,5))

Output:

3D Input array:
[[[0 1]
  [2 3]]

 [[4 5]
  [6 7]]]

Result (neighbor >=5):
[[[False  True]
  [ True  True]]

 [[ True  True]
  [ True  True]]]

This solution efficiently handles arrays of any dimension and size, leveraging NumPy's vectorized operations to avoid slow element-wise loops.

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