New pandas series with flag according to another series

user_12383107 0
Viewed 90

I have a dataframe similar to this:

>>> d = {'ID': ['ID1', 'ID2', 'ID3', 'ID4', 'ID5', 'ID6', 'ID7', 'ID8', 'ID9', 'ID10'], 
         'A': [1, 1, 1, 1, 2, 2, 2, 2, 2, 2], 
         'B': [145,158,240,250,199,204,300,350,467,578]}
>>> df = pd.DataFrame(data=d)

I want to create a new series, F, to flag every 100 units of column B (starting to count from the first value in column B, not from 0). The numbers from column B "restart" for every number in the column A. For a new number in column A, it should start a new flag and take the respective value from column B as first number of the new range of 100. To clarify, the expected outcome for this situation would be:

>>> outcome = {'ID': ['ID1', 'ID2', 'ID3', 'ID4', 'ID5', 'ID6', 'ID7', 'ID8', 'ID9', 'ID10'], 
           'A': [1, 1, 1, 1, 2, 2, 2, 2, 2, 2], 
           'B': [145,158,240,250,199,204,300,350,467,578],
           'F': ['F1','F1','F1','F2','F3','F4','F4','F5','F6','F7']}
>>> outcome
      A    B    F
ID1   1   145   F1
ID2   1   158   F1
ID3   1   240   F1
ID4   1   250   F2
ID5   2   199   F3
ID6   2   204   F3
ID7   2   300   F4
ID8   2   350   F4
ID9   2   467   F5
ID10  2   578   F6

I hope it all made sense, thanks in advance!

pandas python
5 Answers

You can do:

import numpy as np

df['d100'] = df.groupby('A')['B'].diff().fillna(0)
df['d100'] = df.groupby('A')['d100'].cumsum() // 100

df['F'] = np.where(df['A'].ne(df['A'].shift()) | df['d100'].ne(df['d100'].shift()), 1, 0).cumsum()
df['F'] = 'F' + df['F'].astype(str)

df.drop('d100', axis=1, inplace=True)

Outputs:

     ID  A    B   F
0   ID1  1  145  F1
1   ID2  1  158  F1
2   ID3  1  240  F1
3   ID4  1  250  F2
4   ID5  2  199  F3
5   ID6  2  204  F3
6   ID7  2  300  F4
7   ID8  2  350  F4
8   ID9  2  467  F5
9  ID10  2  578  F6
user_116101860

This is a (brute force) solution I would propose:

df = df.reset_index()                # iloc is easier with a clean integer index
B0 = df['B'][0]                      # initialize B

df['F'] = ''                         # create a result column 'F'
df.loc[0,'F'] = 'F1'                 # set the first result
idx = 1                              # initialize your index 
for i in range(1,len(df)):           # iterate over all rows
    if(df['A'][i] == df['A'][i-1]):  # condition 1 : Ai == Ai-1
        if((df['B'][i]-B0)>100):     # condition 2 : Bi - B0 > 100
            idx += 1                 # increment index
            B0 = df.loc[i,'B']       # reset B0
    else:                            # Ai != Ai-1
        idx +=1                      # increment index
        B0 = df.loc[i,'B']           # reset B0

    df.loc[i,'F'] = 'F' + str(idx)   # set output Fi

Interested to see if somebody can provide a more beautiful solution.

user_50083800

A simplification which is shorter but less readable came to my mind and I post it as another answer to let you chose which one you prefer:

df = df.reset_index()                # iloc is easier with a clean integer index
B0 = df['B'][0]                      # initialize B

df['F'] = ''                         # create a result column 'F'
df.loc[0,'F'] = 'F1'                 # set the first result
idx = 1                              # initialize your index 
for i in range(1,len(df)):           # iterate over all rows
    if(df['A'][i] != df['A'][i-1]) |  if((df['B'][i]-B0)>100):     # combining both conditions
        idx += 1                 # increment index
        B0 = df.loc[i,'B']       # reset B0

    df.loc[i,'F'] = 'F' + str(idx)   # set output Fi
user_50083800

I couldn't find a way without a loop, that also worked with slightly different example data. I include the different test data at the end.

First setup the example data

import pandas as pd
import numpy as np

d = {'ID': ['ID1', 'ID2', 'ID3', 'ID4', 'ID5', 'ID6', 'ID7', 'ID8', 'ID9', 'ID10'], 
         'A': [1, 1, 1, 1, 2, 2, 2, 2, 2, 2], 
         'B': [145,158,240,250,199,204,300,350,467,578]}
df = pd.DataFrame(data=d)
print(df)

Out:

     ID  A    B
0   ID1  1  145
1   ID2  1  158
2   ID3  1  240
3   ID4  1  250
4   ID5  2  199
5   ID6  2  204
6   ID7  2  300
7   ID8  2  350
8   ID9  2  467
9  ID10  2  578

Create a helping Series to find the differences in rows and a helper function to find the rows that sum to more than 100.

gr = df.groupby('A')['B'].apply(lambda x: x.diff()).fillna(0)

less100 = np.frompyfunc(lambda x,y: 0 if x + y > 100 else x + y, 2, 1)

df['F'] = 'F' + gr.groupby(df.A).apply(
        lambda x: ~less100.accumulate(x.astype(object)).astype('bool')
    ).cumsum().astype('str')
print(df)

Out:

     ID  A    B   F
0   ID1  1  145  F1
1   ID2  1  158  F1
2   ID3  1  240  F1
3   ID4  1  250  F2
4   ID5  2  199  F3
5   ID6  2  204  F3
6   ID7  2  300  F4
7   ID8  2  350  F4
8   ID9  2  467  F5
9  ID10  2  578  F6

Slightly different data with row 4, column B = 99

d = {'ID': ['ID1', 'ID2', 'ID3', 'ID4', 'ID5', 'ID6', 'ID7', 'ID8', 'ID9', 'ID10'], 
         'A': [1, 1, 1, 1, 2, 2, 2, 2, 2, 2], 
         'B': [145,158,240,250, 99,204,300,350,467,578]}
df = pd.DataFrame(data=d)
print(df)

Out:

     ID  A    B
0   ID1  1  145
1   ID2  1  158
2   ID3  1  240
3   ID4  1  250
4   ID5  2   99
5   ID6  2  204
6   ID7  2  300
7   ID8  2  350
8   ID9  2  467
9  ID10  2  578

gr = df.groupby('A')['B'].apply(lambda x: x.diff()).fillna(0)
less100 = np.frompyfunc(lambda x,y: 0 if x + y > 100 else x + y, 2, 1)
df['F'] = 'F' + gr.groupby(df.A).apply(
        lambda x: ~less100.accumulate(x.astype(object)).astype('bool')
    ).cumsum().astype('str')
print(df)

Out:

     ID  A    B   F
0   ID1  1  145  F1
1   ID2  1  158  F1
2   ID3  1  240  F1
3   ID4  1  250  F2
4   ID5  2   99  F3
5   ID6  2  204  F4
6   ID7  2  300  F4
7   ID8  2  350  F5
8   ID9  2  467  F6
9  ID10  2  578  F7
user_142777220

A group by with head(/minimum) will help, may be faster too.

# Get the first row of each group, it is minimum
df_grp = df.groupby('A').head(1)

# Get the difference of  first row to each row
df_res = pd.merge(df, df_grp[['A','B']], how='inner', on=['A'])

# Make bucket name as new column
l_1 = ['F' + str(x) for x in (df_res['B_x'] - df_res['B_y'])//100 + 1]
df_res['F'] = pd.Series(l_1)

# Drop the unnecessary columns and rename
df_res.drop(columns=['B_y'], inplace=True)
df_res.rename(columns={'B_x':'B'}, inplace=True)

and output

ID  A    B   F
0   ID1  1  145  F1
1   ID2  1  158  F1
2   ID3  1  240  F1
3   ID4  1  250  F2
4   ID5  2  199  F1
5   ID6  2  204  F1
6   ID7  2  300  F2
7   ID8  2  350  F2
8   ID9  2  467  F3
9  ID10  2  578  F4
>
user_75530060
Related