我有以下 pandas 系列(以列表形式表示):
[7,2,0,3,4,2,5,0,3,4]
我想定义一个新的系列,返回到最后一个零的距离。这意味着我想要以下输出:
[1,2,0,1,2,3,4,0,1,2]
如何在pandas中以最有效的方式做到这一点?
复杂度是
O(n)
。会减慢速度的是在 python 中执行 for
循环。如果序列中有 k
个零,并且 log k
与序列的长度相比可以忽略不计,则 O(n log k)
解决方案将是:
>>> izero = np.r_[-1, (ts == 0).nonzero()[0]] # indices of zeros
>>> idx = np.arange(len(ts))
>>> idx - izero[np.searchsorted(izero - 1, idx) - 1]
array([1, 2, 0, 1, 2, 3, 4, 0, 1, 2])
Pandas 中的解决方案有点棘手,但可能看起来像这样(
s
是你的系列):
>>> x = (s != 0).cumsum()
>>> y = x != x.shift()
>>> y.groupby((y != y.shift()).cumsum()).cumsum()
0 1
1 2
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
dtype: int64
最后一步,使用 Pandas 食谱中的“itertools.groupby”配方这里。
一个可能性能不高(尚未真正检查过)但在步骤方面更容易理解的解决方案(至少对我来说)是:
df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})
df
df['flag'] = np.where(df['X'] == 0, 0, 1)
df['cumsum'] = df['flag'].cumsum()
df['offset'] = df['cumsum']
df.loc[df.flag==1, 'offset'] = np.nan
df['offset'] = df['offset'].fillna(method='ffill').fillna(0).astype(int)
df['final'] = df['cumsum'] - df['offset']
df
有时会令人惊讶地看到使用 Cython 获得类似 C 的速度是多么简单。假设您的列的
.values
给出 arr
,那么:
cdef int[:, :, :] arr_view = arr
ret = np.zeros_like(arr)
cdef int[:, :, :] ret_view = ret
cdef int i, zero_count = 0
for i in range(len(ret)):
zero_count = 0 if arr_view[i] == 0 else zero_count + 1
ret_view[i] = zero_count
注意使用类型化内存视图,它的速度非常快。您可以使用
@cython.boundscheck(False)
装饰函数来进一步加快速度。
另一种选择
df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})
zeros = np.r_[-1, np.where(df.X == 0)[0]]
def d0(a):
return np.min(a[a>=0])
df.index.to_series().apply(lambda i: d0(i - zeros))
或者使用纯numpy
df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})
a = np.arange(len(df))[:, None] - np.r_[-1 , np.where(df.X == 0)[0]][None]
np.min(a, where=a>=0, axis=1, initial=len(df))
accumulate
实现此目的的另一种方法。唯一的问题是,要将计数器初始化为零,您需要在系列值前面插入一个零。
import numpy as np
# Define Python function
f = lambda a, b: 0 if b == 0 else a + 1
# Convert to Numpy ufunc
npf = np.frompyfunc(f, 2, 1)
# Apply recursively over series values
x = npf.accumulate(np.r_[0, s.values])[1:]
print(x)
array([1, 2, 0, 1, 2, 3, 4, 0, 1, 2], dtype=object)
这里有一个不使用groupby的方法:
((v:=pd.Series([7,2,0,3,4,2,5,0,3,4]).ne(0))
.cumsum()
.where(v.eq(0)).ffill().fillna(0)
.rsub(v.cumsum())
.astype(int)
.tolist())
输出:
[1, 2, 0, 1, 2, 3, 4, 0, 1, 2]
也许 pandas 不是最好的工具,正如 @behzad.nouri 的答案,但是这里有另一种变体:
df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})
z = df.ne(0).X
z.groupby((z != z.shift()).cumsum()).cumsum()
0 1
1 2
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
Name: X, dtype: int64
解决方案2:
如果您编写以下代码,您将获得几乎您需要的所有内容,除了第一行从 0 而不是 1 开始:
df = pd.DataFrame({'X': [7, 2, 0, 3, 4, 2, 5, 0, 3, 4]})
df.eq(0).cumsum().groupby('X').cumcount()
0 0
1 1
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
dtype: int64
发生这种情况是因为累加和从 0 开始计数。为了得到想要的结果,我在第一行添加了 0,计算了所有内容,然后在末尾去掉 0,得到:
x = pd.Series([0], index=[0])
df = pd.concat([x, df])
df.eq(0).cumsum().groupby('X').cumcount().reset_index(drop=True).drop(0).reset_index(drop=True)
0 1
1 2
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
dtype: int64
s1=pd.Series([7,2,0,3,4,2,5,0,3,4])
k=(s1==0).idxmax() #get the idx of the first occurrence of 0
s1=(s1!=0).cumsum()
0 0
1 0
2 1
3 1
4 1
5 1
6 1
7 2
8 2
9 2
dtype: int64
grouped = s1.groupby(s1)
grouped=grouped.transform(lambda x:np.arange(x.size))
0 0
1 1
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
grouped[:k]=grouped[:k]+1
0 1
1 2
2 0
3 1
4 2
5 3
6 4
7 0
8 1
9 2
dtype: int64