NumPy技巧：如何获取一维扁平化数组中指定元素的相邻元素？

问题背景 #

我们有一个已扁平化的一维ndarray数组，希望通过NumPy的向量化API替代循环，高效获取数组中与指定元素（或元素集合）相邻的元素。以下是用户提供的低效循环实现：

import numpy as np

def find_adjacent(v,a,side={'L','R','B'}):
    # v 是扁平化后的一维数组
    # a 是需要相邻的目标元素（或一维元素集合）
    # side 可选值：'L'（仅找左侧相邻a的元素）、'R'（仅找右侧相邻a的元素）、'B'（找左右都相邻a的元素）

    result = np.array([])
    for i in range(v.size):
        if (side == 'L' or side == 'B') and i < a.size:
            continue
        elif (side == 'R' or side == 'B') and v.size-i < a.size:
            continue
        elif side == 'L':
            if np.array_equal(v[i-a.size:i], a):
                np.append(result,v[i])
        elif side == 'R':
            if np.array_equal(v[i+1:i+a.size], a):
                np.append(result,v[i])
        elif side == 'B':
            if np.array_equal(v[i+1:i+a.size], a) and np.array_equal(v[i-a.size:i], a):
                np.append(result,v[i])

    return result

需要解决的问题：

1. 如何用无循环的NumPy风格代码实现相同功能？
2. 若目标元素a是标量，实现是否会更简单？

解决方案 #

一、当a是一维数组（元素集合）时的实现 #

利用NumPy的sliding_window_view（需NumPy 1.20及以上版本）生成滑动窗口，通过向量化比较快速定位匹配位置，彻底避免循环：

import numpy as np

def find_adjacent_vectorized(v, a, side='B'):
    v = np.asarray(v)
    a = np.asarray(a)
    len_a = a.size
    len_v = v.size
    result = np.array([], dtype=v.dtype)

    if side in ('L', 'B'):
        # 处理左侧匹配：找前len_a个元素等于a的位置
        if len_a <= len_v:
            left_windows = np.lib.stride_tricks.sliding_window_view(v, len_a)
            left_matches = np.all(left_windows == a, axis=1)
            # 匹配窗口的结束位置即为目标元素索引
            left_indices = np.where(left_matches)[0] + len_a
            result = np.concatenate([result, v[left_indices]])

    if side in ('R', 'B'):
        # 处理右侧匹配：找后len_a个元素等于a的位置
        if len_a <= len_v:
            right_windows = np.lib.stride_tricks.sliding_window_view(v, len_a)
            right_matches = np.all(right_windows == a, axis=1)
            # 匹配窗口的起始位置前一位即为目标元素索引，过滤边界无效值
            right_indices = np.where(right_matches)[0] - 1
            right_indices = right_indices[right_indices >= 0]
            result = np.concatenate([result, v[right_indices]])

    if side == 'B':
        # 取左右匹配的交集，仅保留同时满足两侧条件的元素
        left_set = set(left_indices)
        right_set = set(right_indices)
        common_indices = np.array(list(left_set & right_set))
        result = v[common_indices]

    # 去重避免重复元素（与原循环逻辑保持一致）
    return np.unique(result) if side == 'B' else result

二、当a是标量时的实现 #

标量场景下逻辑更简洁，无需处理多元素窗口匹配，直接通过相邻位置的元素比较生成掩码即可：

def find_adjacent_scalar(v, a, side='B'):
    v = np.asarray(v)
    len_v = v.size
    result = np.array([], dtype=v.dtype)

    if side in ('L', 'B'):
        # 左侧相邻标量a：元素v[i]的前一个元素等于a
        left_mask = np.zeros(len_v, dtype=bool)
        left_mask[1:] = v[:-1] == a
        result = np.concatenate([result, v[left_mask]])

    if side in ('R', 'B'):
        # 右侧相邻标量a：元素v[i]的后一个元素等于a
        right_mask = np.zeros(len_v, dtype=bool)
        right_mask[:-1] = v[1:] == a
        result = np.concatenate([result, v[right_mask]])

    if side == 'B':
        # 同时满足左右两侧相邻条件
        both_mask = left_mask & right_mask
        result = v[both_mask]

    return np.unique(result) if side == 'B' else result

说明 #

• 滑动窗口方法sliding_window_view是NumPy官方推荐的向量化序列匹配方案，相比循环效率提升显著，尤其适合大规模数组场景。
• 标量场景的实现确实更简单，省去了多元素窗口的匹配逻辑，直接通过单元素比较即可完成。
• 代码完全对齐原循环的逻辑，包括边界处理、side='B'的双向匹配要求等。

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