Numpy fastest way to get indexes of neighbors with current value (flood fill)-CodePudding

I need to find fast way to get indicies of neighbors with values like current

For example:

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

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

6th element has value 1, so I need full slice containing 6th and all it's neighbors with same value

It is like a part of flood fill algorithm. Is where a way to do it fast in numpy? Is where a way for 2D array?

CodePudding user response：

I don't know how to solve this problem with numpy but If you use pandas, you might get the result that you want with this:

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)]

Output:

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

CodePudding user response：

Here is a numpy solution. I think you can improve it by a little more work:

def func(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 np.arange(start, end)

sample output:

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

This would specially be faster if you have few changes in your arr (relative to array size) and you are looking for multiple of those index searches in the same array. Otherwise, faster solutions come to mind.