python - Pandas 数据帧中特定列后给定大小的第一个零序列的长度

Question

假设我有一个这样的数据框:

        ID      0   1   2   3   4   5   6   7   8   ... 81  82  83  84  85  86  87  88  89  90  total  day_90
-------------------------------------------------------------------------------------------------------------
0       A       2   21  0   18  3   0   0   0   2   ... 0   0   0   0   0   0   0   0   0   0    156   47
1       B       0   20  12  2   0   8   14  23  0   ... 0   0   0   0   0   0   0   0   0   0    231   35
2       C       0   38  19  3   1   3   3   7   1   ... 0   0   0   0   0   0   0   0   0   0     78   16
3       D       3   0   0   1   0   0   0   0   0   ... 0   0   0   0   0   0   0   0   0   0      5   3

其中最后一列 [day_90] 包含哪一列 ([0] - [90]) 累积每行 [total] 的 90% 的值。澄清一下，以第一行为例:在第47列中，ID A总共命中了他将在90天内完成的156个事件中的90%。

我需要的是:对于每一行，计算大于 7(或任何预定义的任意数字)的第一个 0 序列的长度。因此，例如:对于第一行，我想知道第 47 列之后的第一个零序列有多长，但前提是该序列连续超过 7 个零。如果有 6 个零，然后有一个非零，那我就不想数了。

最后，我想将此结果存储在 [day_90] 之后的新列中。因此，如果 ID A 在第 47 列之后有一个由 10 个零组成的序列，我想添加一个新列 [0_sequence] 来保存该 ID 的值 10。

我真的不知道从哪里开始。感谢任何帮助 =)

Answer 1

您的问题基本上是孤岛和缺口问题的变体:非零会创建一个新的“孤岛”，而 0 会扩展当前的孤岛。并且您想找到第一个具有一定大小的岛屿。在我回答您的问题之前，让我向您介绍一下该问题的缩小版本。

假设您有一个系列:

>>> a = pd.Series([0,0,0,13,0,0,4,12,0,0])
0     0
1     0
2     0
3    13
4     0
5     0
6     4
7    12
8     0
9     0

并且您想找到长度至少为 3 个元素的第一个 0 序列的长度。让我们首先将它们分配到“孤岛”中:

# Every time the number is non-zero, a new "island" is created
>>> b = (a != 0).cumsum()
0    0  <-- island 0
1    0
2    0
3    1  <-- island 1
4    1
5    1
6    2  <-- island 2
7    3  <-- island 3
8    3
9    3

对于每个岛屿，我们只对等于 0 的元素感兴趣:

>>> c = b[a == 0]
0    0
1    0
2    0
4    1
5    1
8    3
9    3

现在让我们确定每个岛的大小:

>>> d = c.groupby(c).count()
0    3  <-- island 0 is of size 3
1    2  <-- island 1 is of size 2
3    2  <-- island 3 is of size 2
dtype: int64

并过滤大小 >= 3 的岛屿:

>>> e = d[d >= 3]
0    3

如果 e 不为空，则答案是 e 的第一个元素(岛 0，大小 3)。否则，没有岛屿符合我们的标准。

---

第一次尝试

并将其应用于您的问题:

def count_sequence_length(row, n):
    """Return of the length of the first sequence of 0
    after the column in `day_90` whose length is >= n
    """
    if row['day_90'] + n > 90:
        return 0
    
    # The columns after `day_90`
    idx = np.arange(row['day_90']+1, 91)

    a = row[idx]
    b = (a != 0).cumsum()
    c = b[a == 0]
    d = c.groupby(c).count()
    e = d[d >= n]
    
    return 0 if len(e) == 0 else e.iloc[0]

df['0_sequence'] = df.apply(count_sequence_length, n=7, axis=1)

---

第二次尝试

上面的版本很好，但是很慢，因为它计算了所有岛屿的大小。由于您只关心第一个满足标准的岛屿的大小，因此一个简单的 for 循环工作得更快:

def count_sequence_length_2(row, n):
    if row['day_90'] + n > 90:
        return 0
    
    size = 0
    for i in range(row['day_90']+1, 91):
        if row[i] == 0:
            # increase the size of the current island
            size += 1
        elif size >= n:
            # found the island we want. Search no more
            break
        else:
            # create a new island
            size = 0
    return size if size >= n else 0

df['0_sequence'] = df.apply(count_sequence_length_2, n=7, axis=1)

当我对其进行基准测试时，这实现了 10 到 20 倍的速度提升。