Pandas 数据框列上带有计数器的矢量化函数

Question

Consider this pandas dataframe where the condition column is 1 when value is below 5 (any threshold).

import pandas as pd
d = {'value': [30,100,4,0,80,0,1,4,70,70],'condition':[0,0,1,1,0,1,1,1,0,0]}
df = pd.DataFrame(data=d)
df

Out[1]:
   value  condition
0     30          0
1    100          0
2      4          1
3      0          1
4     80          0
5      0          1
6      1          1
7      4          1
8     70          0
9     70          0

What I want is to have all consecutive values below 5 to have the same id and all values above five have 0 (or NA or a negative value, doesn't matter, they just need to be the same). I want to create a new column called new_id that contains these cumulative ids as follows:

   value  condition  new_id
0     30          0       0
1    100          0       0
2      4          1       1
3      0          1       1
4     80          0       0
5      0          1       2
6      1          1       2
7      4          1       2
8     70          0       0
9     70          0       0

In a very inefficient for loop I would do this (which works):

for i in range(0,df.shape[0]):
    if (df.loc[df.index[i],'condition'] == 1) & (df.loc[df.index[i-1],'condition']==0):
        new_id = counter # assign new id 
        counter += 1 

    elif (df.loc[df.index[i],'condition']==1) & (df.loc[df.index[i-1],'condition']!=0):
        new_id = counter-1 # assign current id

    elif (df.loc[df.index[i],'condition']==0):
        new_id = df.loc[df.index[i],'condition'] # assign 0

    df.loc[df.index[i],'new_id'] = new_id
df
  

But this is very inefficient and I have a very big dataset. Therefore I tried different kinds of vectorization but I so far failed to keep it from counting up inside each "cluster" of consecutive points:

# First try using cumsum():
df['new_id'] = 0
df['new_id_temp'] = ((df['condition'] == 1)).astype(int).cumsum()
df.loc[(df['condition'] == 1), 'new_id'] = df['new_id_temp']
df[['value', 'condition', 'new_id']]

# Another try using list comprehension but this just does +1:
[row+1 for ind, row in enumerate(df['condition']) if (row != row-1)]

I also tried using apply() with a custom if else function but it seems like this does not allow me to use a counter.

There is already a ton of similar posts about this but none of them keep the same id for consecutive rows.

Example posts are: Maintain count in python list comprehension Pandas cumsum on a separate column condition Python - keeping counter inside list comprehension python pandas conditional cumulative sum Conditional count of cumulative sum Dataframe - Loop through columns

Answer 1

Welcome to SO! Why not just rely on base Python for this?

def counter_func(l):
    new_id = [0]   # First value is zero in any case
    counter = 0
    for i in range(1, len(l)):
        if l[i] == 0:
            new_id.append(0)
        elif l[i] == 1 and l[i-1] == 0:
            counter += 1
            new_id.append(counter)
        elif l[i] == l[i-1] == 1:
            new_id.append(counter)
        else: new_id.append(None)
    return new_id

df["new_id"] = counter_func(df["condition"])

Looks like this

   value  condition  new_id
0     30          0       0
1    100          0       0
2      4          1       1
3      0          1       1
4     80          0       0
5      0          1       2
6      1          1       2
7      4          1       2
8     70          0       0
9     70          0       0

Edit :

You can also use numba , which sped up the function quite a lot for me about : about 1sec to ~60ms.

You should input numpy arrays into the function to use it, meaning you'll have to df["condition"].values .

from numba import njit
import numpy as np
@njit
def func(arr):
    res = np.empty(arr.shape[0])
    counter = 0
    res[0] = 0 # First value is zero anyway
    for i in range(1, arr.shape[0]):
        if arr[i] == 0:
            res[i] = 0
        elif arr[i] and arr[i-1] == 0:
            counter += 1
            res[i] = counter
        elif arr[i] == arr[i-1] == 1:
            res[i] = counter
        else: res[i] = np.nan
    return res

df["new_id"] = func(df["condition"].values)

Answer 2

You can use the cumsum() , as you did in your first try, just modify it a bit:

# calculate delta
df['delta'] = df['condition']-df['condition'].shift(1)
# get rid of -1 for the cumsum (replace it by 0)
df['delta'] = df['delta'].replace(-1,0)

# cumulative sum conditional: multiply with condition column
df['cumsum_x'] = df['delta'].cumsum()*df['condition']