Python 循环遍历数据帧行，直到第一次满足条件

Question

我有一个 Pandas 数据框，我想在其中循环遍历其行并计算从第一行到第二行的度量，如果在那里找不到，请检查从第一行到第三行、第四行等，并将该度量与另一个值进行比较。 我想获得第一次满足条件的行号。 举一个具体的例子，对于长度为 30 的数据帧，它可能来自df.iloc[0:10] df.iloc[10:15]和df.iloc[15:27] , df.iloc[27:30] ，其中值 10、15、27 存储在列表中。

一个示例数据框：

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randint(0,100, size=(100, 1)), columns=list('A'))
df  
    A
0   5
1  11
2   8
3   1
4  16
5  24

some_value = 20 
mylist = []
for i in range(len(df)):
    for j in range(i+2, range(len(df)):
        # Metric calculated on the relevant rows
        metric = df.iloc[i:j]['A'].sum()
        if metric >= some_value:
           mylist.append(j)
           break

循环从df.iloc[0:2] ，计算 5+11，因为它不大于 some_value (20)，它传递给df.iloc[0:3] 。 这一次，由于 5+11+8 大于 some_value，我想保存这个数字 (2) 并且不检查df.iloc[0:4] 。 然后循环应该再次从df.iloc[3:5]开始检查（1+16），因为不满足条件，继续df.iloc[3:6] （1+16+24）和以此类推，并在满足条件时保存积分。

这种情况下的示例输出是一个包含值的列表： [2, 5]

我写了上面的代码，但不能完全实现我想要的。你能帮忙解决这个问题吗？ 谢谢。

Answer 1

目前，您的循环是 O(n^2)。 但是一旦找到与 i 的起始值匹配的值，您的外循环必须从 i+1 重新开始，而您不想从那里开始。 你想从 j 开始。 这是对您的代码的快速修复。

我目前没有 numpy，所以我使用 python 列表作为数据。

data = [5, 11, 8, 1, 16, 24]
some_value = 20 
mylist = []
j = 0
for i in range(len(data)):
    # can't change iteration so just skip ahead with continue
    if i < j:
        continue
    # range expects second argument to be past the end
    # dunno if df is the same, but probably?
    for j in range(i+1, len(data)+1):
        metric = sum(data[i:j])
        if metric >= some_value:
            mylist.append(j-1)
            break
print(mylist)

[2, 5]

我建议在一个循环中执行此操作，并保持运行总数（累加器）。 在这里，我有点喜欢返回范围，以防您想拼接 df：

data = [5, 11, 8, 1, 16, 24]
threshold = 20

def accumulate_to_threshold(data, threshold):
    start = 0
    total = 0
    for index, item in enumerate(data):
        total += item
        if total > threshold:
            yield (start, index+1)
            total = 0
            start = index+1
    # leftovers below threshold here

for start, end in accumulate_to_threshold(data, threshold):
    sublist = data[start:end]
    print (sublist, "totals to", sum(sublist))

[5, 11, 8] 总数为 24
[1, 16, 24] 总计 41

当然，您可以生成索引并从上面获取 [2, 5]，而不是生成一个范围。

Answer 2

我的方法是：

• numpy.reshape(values, newshape, ...)
• .sum(axis=1)
• 布尔掩码

我不知道这是否会以您想要的方式回答您的问题，但我将展示我的大脑如何使用 pandas/numpy 的内置矢量化来处理它，简而言之，循环很麻烦（慢），如果可能：

import pandas as pd
import numpy as np

# made it smaller
df = pd.DataFrame(np.random.randint(0,25, size=(20, 1)), columns=list('A'))

numpy.reshape()和sum()

我们将重塑 col A ，它将值并排移动，然后求和穿过axis=1 ：

将df与下面的re_shaped进行比较。 注意这些值是如何重新排列的

re_shaped = np.reshape(df.A.values, (10, 2))
print(df)

     A
0    5
1   11
2    8
3   23
...
16   6
17  14
18   3
19   0

print(re_shaped)

array([[ 5, 11],
       [ 8, 23],
       ...
       [ 6, 14],
       [ 3,  0]])

summed = re_shaped.sum(axis=1)
print(summed)

array([16, 31, 15, 19, 13, 21, 28, 30, 20,  3])

布尔掩码

some_value = 20
greater_than_some_value = summed[summed >= some_value]
print(greater_than_some_value)

array([31, 21, 28, 30, 20])

你有它。 希望它有所帮助。

Answer 3

您是否考虑过仅使用一个循环：

import pandas as pd
import numpy as np

n = int(1e6)
df = pd.DataFrame({"A": np.random.randint(100, size=n)})

threshold = 20
my_list = []
s = 0
for i, k in enumerate(df["A"].values):
    if s + k > threshold:
        my_list.append(i)
        s = 0
    else:
        s += k

您最终可以使用numba但我认为最好的想法是在您的df使用 reset 计算 cumsum 。

努巴

前一个可以写成一个函数

def fun(vec, threshold=20):
    my_list = []
    s = 0
    for i, k in enumerate(vec):
        if s + k > threshold:
            my_list.append(i)
            s = 0
        else:
            s += k
    return my_list

我们可以使用 numba

from numba import jit

@jit(nopython=True, cache=True, nogil=True)
def fun_numba(vec, threshold=20):
    my_list = []
    s = 0
    for i, k in enumerate(vec):
        if s + k > threshold:
            my_list.append(i)
            s = 0
        else:
            s += k
    return my_list

%%timeit -n 5 -r 5
my_list = fun(df["A"].values)

606 ms ± 28 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

%%timeit -n 5 -r 5
my_list = fun_numba(df["A"].values)

59.6 ms ± 20.4 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

这是大约 10 倍的加速。