我用pandas 写了个sumbars，二次for循环，感觉效率很慢，需要用向量化优化一下，麻烦各位看看
def sumbars(df, s, n):
#s表示列名
start = time()
df2 = df[s].copy()
# df2.loc[0]=1000
df['sumbars'] = 0
for i in range(len(df)):
df3 = df2[:i + 1]
l = len(df3)
for j in range(l):
# if i>=8:
# print()
df4 = df3[l - j - 1:]
if df4.sum() >= n:
df.loc[i, 'sumbars'] = j + 1
# df5=df[['trade_date',s,'sumbars']]
break
stop = time()
global sumbarstime
sumbarstime = sumbarstime + stop - start
return df['sumbars']
题主，你的意思是希望找到第s列，以第i行为结尾，满足第j行开始，和大于n的那一个j吗？
如果是这样的话，算法运行效率比较慢的原因是，在if df4.sum() >= n的时候，每次计算完sum(0:i+1)之后，又重新计算一次sum(1:i+1)，再计算sum(2:i+1)，再计算sum(3:i+1)，没有充分利用上一次得到的计算结果。这样的运算之后需要计算1+2+3+...+i+1,是O(n方)级别的。在加上外层的循环，总的时间复杂度是O(n^3).
你可以计算第一次sum(0:i+1)之后，每次减去a[0],a[1],a[2]。。。这样总的时间复杂度就降低为O(n^2)级别的了。
def sumbars(df,s,n):
#s表示列名
start = time()
sdata = np.array(df[s])
sdata = np.flipud(sdata)
l = len(sdata)
sumbars = np.zeros(l)
for i in range(l):
cumsum = np.cumsum(sdata[i:])
k = np.argmax(cumsum>=n)
if k != 0:
sumbars[l-i-1] = k+1
stop = time()
print(stop-start)
global sumbarstime
sumbarstime = sumbarstime + stop - start
sumbars.astype(int)
df['sumbars'] = sumbars
# df.to_csv("my.csv")
# print(df['sumbars'])
return df['sumbars']更进一步，这个allSum也可以不用每个i都计算一次。只需要每次增加就行。
def sumbars(df, s, n):
#s表示列名
start = time()
df['sumbars'] = 0
allSum = 0
for i in range(len(df)):
allSum += df[s].iloc[i]
eachSum = allSum
for j in range(i+1):
if eachSum >= n:
df.loc[i, 'sumbars'] = j + 1
break
eachSum -= df[s].iloc[j]
stop = time()
global sumbarstime
sumbarstime = sumbarstime + stop - start
return df['sumbars']另外，导致题主之前代码运行比较慢的原因，也可能是多次对df进行拷贝导致的。当DataFrame较大的时候，拷贝也会占用一定时间。
如果我对题主代码的意思理解有偏差，也请评论或私信告诉我。
========
和题主沟通后，题主的意思是找到最后一个符合大于等于n的行下标，同时使用numpy进行优化。优化后的代码如下。
def sumbars(df, s, n):
start = time()
sdata = np.array(df[s])
sdata = np.flipud(sdata)
l = len(sdata)
sumbars = np.zeros(l)
for i in range(l):
cumsum = np.cumsum(sdata[i:])
k = np.argmax(cumsum>=n)
if k != 0:
sumbars[l-i-1] = k+1
stop = time()
print(stop-start)
global sumbarstime
sumbarstime = sumbarstime + stop - start
sumbars.astype(int)
df['sumbars'] = sumbars
# df.to_csv("my.csv")
# print(df['sumbars'])
return df['sumbars']
把函数的输入s改为series后，用JIT：
from numba import njit, prange, int64
@njit(int64[:](int64[:], int64), parallel=True)
def sumbarsKernel(sdata, n):
l = len(sdata)
fullCumsum = np.cumsum(sdata)
sumbars = np.zeros(l, dtype=np.int64)
nShared = np.int64(n)
for i in prange(l):
n1 = l - i - 1
m = fullCumsum[i - 1] if i > 0 else 0
k = np.searchsorted(fullCumsum[i:], n + m)
if k == len(fullCumsum) - i:
k = 0
if (k != 0) or ((fullCumsum[i + k] != m) and (nShared == 1)):
sumbars[l - i - 1] = k + 1
return sumbars
def sumbars(s, n):
start = time()
sdata = np.array(s, dtype=np.int64)
sdata = np.flipud(sdata)
sumbars = sumbarsKernel(sdata, n)
sumbars = sumbars.astype(int)
stop = time()
global sumbarstime
sumbarstime += stop - start
return sumbars