Home > Back-end >  pandas: pivot on two columns
pandas: pivot on two columns

Time:10-02

I have the following data:

import pandas as pd, numpy as np
dates = pd.date_range('01/01/2022', '01/11/2022', freq = 'D')
values = [0,0,1,1,0,0,1,1,1,0,1]
df = pd.DataFrame({'date': dates, 'value': values})
df

    date    value
0   2022-01-01  0
1   2022-01-02  0
2   2022-01-03  1
3   2022-01-04  1
4   2022-01-05  0
5   2022-01-06  0
6   2022-01-07  1
7   2022-01-08  1
8   2022-01-09  1
9   2022-01-10  0
10  2022-01-11  1

I would like to transform this such that i end up with a 'start' and 'end' column such that start is the first occurrence of 1, and end is the last consecutive occurrence of 1. Basically i should end up with this:

start      end
2022-01-03 2022-01-04  
2022-01-07 2022-01-09
2022-01-11

So what i have done so far is the following:

conditions = [
    (df.value == 1) & (df.value.shift(1) == 0),
    (df.value == 1) & (df.value.shift(-1) == 0)]
choices = ['start', 'end']
df['value'] =  np.select(conditions, choices, default=pd.NA)
df = df.dropna()
df.pivot(columns='value')

    date
value   end         start
2       NaT         2022-01-03
3       2022-01-04  NaT
6       NaT         2022-01-07
8       2022-01-09  NaT
10      NaT         2022-01-11

As you can see it is almost there..and i could do some addition fiddling to get what i want at this point - but i feel maybe im approaching this the wrong way.
Is there a better, more efficient way to solve this problem?

CodePudding user response:

I would use a groupby.agg here:

# which rows have value 1?
m = df['value'].eq(1)

(df[m] # keep only value==1
 .groupby(m.ne(m.shift()).cumsum()) # group by consecutive values
 ['date'].agg(['first', 'last'])    # get first and last date
 .reset_index(drop=True)
)

output:

       first       last
0 2022-01-03 2022-01-04
1 2022-01-07 2022-01-09
2 2022-01-11 2022-01-11

CodePudding user response:

split data frame into chunks and then extract the first / last date from the 1-chunks:

splits = np.split(df, np.flatnonzero(df['value'].diff() != 0))
pd.DataFrame([
    (split['date'].iat[0], split['date'].iat[-1]) 
    for split in splits 
    if len(split) > 0 and split['value'].iat[0] == 1
], columns=['start', 'end'])

#       start        end
#0 2022-01-03 2022-01-04
#1 2022-01-07 2022-01-09
#2 2022-01-11 2022-01-11
  • Related