python - 有效地查找 pandas 数据框中多行的形状

Question

这是关于在 pandas 数据框中有效查找“2D 形状”。我们生成了一个如下所示的框架:

l = ['>','<','v','^']

def gendf(r,c): 
    return pd.DataFrame(np.random.choice(l,size=(r,c)),columns=[i for i in range(c)])

我正在寻找以下形状:

    ^
<   ?   >
    v

用“？”是任何角色。我有一个非常低效的嵌套循环，它缓慢地完成这项工作:

def co(df):
    c=[]
    for i in list(df.index)[:-2]:
        for j in list(df.columns)[1:-1]:
            if df.loc[i,j]=='^':
                if df.loc[i+2,j]=='v' and df.loc[i+1,j-1]=='<' and df.loc[i+1,j+1]=='>':
                    c.append((i,j))
    return c

它只是查找“顶部”，并在找到时检查其余部分。作为一个嵌套循环，它非常慢——有什么函数可以让我更有效地做到这一点吗？

编辑:给定方法的计时器

%timeit co(df)
8.08 s ± 260 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit x1(df) #Chris A's method
1.89 s ± 102 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit x2(df) #sampers' method
6.22 s ± 289 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果有更好的方法，请告诉我 - 这已经很有帮助了!

Answer 1

IIUC，你可以试试numpy.argwhere使用 DataFrame.shift 具有多个 bool 条件和 DataFrame.eq :

# setup
np.random.seed(0)
df = gendf(100, 100)

cond1 = df.eq('<').shift(axis=1)
cond2 = df.eq('>').shift(-1, axis=1)
cond3 = df.eq('^').shift()
cond4 = df.eq('v').shift(-1)

coords = np.argwhere((cond1 & cond2 & cond3 & cond4).to_numpy())

[输出]

[[ 1  1]
 [ 3 85]
 [11 16]
 [14 87]
 [22 92]
 [24 25]
 [25 14]
 [26 38]
 [27 69]
 [27 85]
 [32 54]
 [36 61]
 [39 29]
 ...
 [69 82]
 [70 76]
 [76 41]
 [79 81]
 [81 60]
 [82 13]
 [83 83]
 [87 12]
 [89 16]
 [94  5]
 [94 69]
 [98 16]]