python - 在 pandas 数据框中查找到下一个更高值的距离

Question

我有一个包含浮点值的数据框

my_df = pd.DataFrame([1,2,1,4,3,2,5,4,7])

我试图找到每个数字，当(需要向前移动多少个索引)直到找到比当前数字大的下一个数字，如果没有更大的数字，我用一些值标记它(例如 999999)。

所以对于上面的例子，正确答案应该是

result = [1,2,1,3,2,1,2,1,999999]

目前我已经通过使用 itertuples 的非常慢的双循环解决了这个问题(意味着 O(n^2))

有更聪明的方法吗？

Answer 1

这是一个基于 numpy 的利用广播的方案:

a = my_df.squeeze().to_numpy() # my_df.squeeze().values for versions 0.24.0.<

diff_mat = a - a[:,None]
result = (np.triu(diff_mat)>0).argmax(1) - np.arange(diff_mat.shape[1])
result[result <= 0] = 99999

<小时/>

print(result)

array([    1,     2,     1,     3,     2,     1,     2,     1, 99999],
      dtype=int64)

<小时/>

其中 diff_mat 是距离矩阵，我们正在查找从主对角线开始的值，这些值大于 0:

array([[ 0,  1,  0,  3,  2,  1,  4,  3,  6],
       [-1,  0, -1,  2,  1,  0,  3,  2,  5],
       [ 0,  1,  0,  3,  2,  1,  4,  3,  6],
       [-3, -2, -3,  0, -1, -2,  1,  0,  3],
       [-2, -1, -2,  1,  0, -1,  2,  1,  4],
       [-1,  0, -1,  2,  1,  0,  3,  2,  5],
       [-4, -3, -4, -1, -2, -3,  0, -1,  2],
       [-3, -2, -3,  0, -1, -2,  1,  0,  3],
       [-6, -5, -6, -3, -4, -5, -2, -3,  0]], dtype=int64)

我们有 np.triu 来实现:

np.triu(diff_mat)

array([[ 0,  1,  0,  3,  2,  1,  4,  3,  6],
       [ 0,  0, -1,  2,  1,  0,  3,  2,  5],
       [ 0,  0,  0,  3,  2,  1,  4,  3,  6],
       [ 0,  0,  0,  0, -1, -2,  1,  0,  3],
       [ 0,  0,  0,  0,  0, -1,  2,  1,  4],
       [ 0,  0,  0,  0,  0,  0,  3,  2,  5],
       [ 0,  0,  0,  0,  0,  0,  0, -1,  2],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  3],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0]], dtype=int64)

通过检查哪些值大于 0，并获取 bool ndarray 的 argmax，我们将找到第一个大于 0 的值每行 0:

(np.triu(diff_mat)>0).argmax(1)
array([1, 3, 3, 6, 6, 6, 8, 8, 0], dtype=int64)

我们只需要减去主对角线到开头对应的偏移即可