python - Numpy 获取具有当前值的邻居索引的最快方法(洪水填充)

Question

我需要找到快速的方法来获取具有当前值的邻居的索引

例如:

arr = [0, 0, 0, 1, 0, 1, 1, 1, 1, 0]

indicies = func(arr, 6)
# [5, 6, 7, 8]

第 6 个元素的值为 1，因此我需要包含第 6 个元素及其所有具有相同值的邻居的完整切片

它就像洪水填充算法的一部分。有没有办法在 numpy 中快速完成？ 有没有办法实现二维数组？

编辑

让我们看看一些性能测试:

import numpy as np
import random

np.random.seed(1488)

arr = np.zeros(5000)
for x in np.random.randint(0, 5000, size = 100):
    arr[x:x+50] = 1

我将比较@Ehsan 的函数:

def func_Ehsan(arr, idx):
    change = np.insert(np.flatnonzero(np.diff(arr)), 0, -1)
    loc = np.searchsorted(change, idx)
    start = change[max(loc-1,0)]+1 if loc<len(change) else change[loc-1]
    end = change[min(loc, len(change)-1)]
    return (start, end)

change = np.insert(np.flatnonzero(np.diff(arr)), 0, -1)
def func_Ehsan_same_arr(arr, idx):
    loc = np.searchsorted(change, idx)
    start = change[max(loc-1,0)]+1 if loc<len(change) else change[loc-1]
    end = change[min(loc, len(change)-1)]
    return (start, end)

用我的纯Python函数:

def my_func(arr, index):
    
    val = arr[index]
    size = arr.size
    
    end = index + 1
    while end < size and arr[end] == val:
        end += 1
    start = index - 1
    while start > -1 and arr[start] == val:
        start -= 1
        
    return start + 1, end

看一下:

np.random.seed(1488)
%timeit my_func(arr, np.random.randint(0, 5000))
# 42.4 µs ± 700 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)


np.random.seed(1488)
%timeit func_Ehsan(arr, np.random.randint(0, 5000))
# 115 µs ± 1.92 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

np.random.seed(1488)
%timeit func_Ehsan_same_arr(arr, np.random.randint(0, 5000))
# 18.1 µs ± 953 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

有没有办法通过 numpy 使用相同的逻辑，而无需 C 模块/Cython/Numba/python 循环？并使其更快!

Answer 1

我不知道如何用 numpy 解决这个问题，但是如果你使用 pandas，你可能会得到你想要的结果:

import pandas as pd
df=pd.DataFrame(arr,columns=["data"])
df["new"] = df["data"].diff().ne(0).cumsum()
[{i[0]:j.index.tolist()} for i,j in df.groupby(["data","new"],sort=False)]

输出:

[{0: [0, 1, 2]}, {1: [3]}, {0: [4]}, {1: [5, 6, 7, 8]}, {0: [9]}]