>我有一个包含浮点值的数据框
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((
有没有更聪明的方法可以做到这一点?
这是一个利用广播的基于 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
,并取布尔ndarray
的argmax
,我们将在每行中找到大于0
的第一个值:
(np.triu(diff_mat)>0).argmax(1)
array([1, 3, 3, 6, 6, 6, 8, 8, 0], dtype=int64)
我们只需要减去从主对角线到开头的相应偏移
量