Create new variable based on groupby & value comparison

Question

I apologize for the not quite correct name of this post. The task consists of several parts, I do not know how to formulate exactly. I have a dataframe like this

df = pd.DataFrame({'id': (1,1,1,2,2), 'begin': ('01.01.2018','01.02.2018', '01.03.2018', '01.01.2018', '01.02.2018'),
    'end': ('01.02.2018','01.03.2018', '01.04.2018', '01.02.2018', '01.03.2018')})
df['begin']= pd.to_datetime(df['begin'])
df['end']= pd.to_datetime(df['end'])
df
      begin        end     id
0   2018-01-01  2018-01-02  1
1   2018-01-02  2018-01-03  1
2   2018-01-03  2018-01-04  1
3   2018-01-01  2018-01-02  2
4   2018-01-02  2018-01-03  2

where begin is a project start time, end is a project end time and id is an employee. I want to create a new variable

      begin       end      id  new
0   2018-01-01  2018-01-02  1   3
1   2018-01-02  2018-01-03  1   2
2   2018-01-03  2018-01-04  1   1
3   2018-01-01  2018-01-02  2   2
4   2018-01-02  2018-01-03  2   1

The new indicates the number of unclosed projects on the date for a specific employee. In row 0 the value is 3 so far as the project start time is 2018-01-01 and this employee has three current projects on this date.

How I solved this problem First I need to get a list of project end / start dates for each employee

gr_end = df.groupby(['id'])['end'].apply(list)
gr_begin = df.groupby(['id'])['begin'].apply(list)

And now the most difficult. I have to compare the project start date for each employee with the end date of each project. I have to calculate all cases where the start date is less than the end date

l = []
for i in gr_begin.index:
    for x in gr_begin[i]:
        cv = [j == 1 for j in gr_end[i] if j > x].count(False)
        l.append(cv)
df['new'] = l

First of all, I ask you to check whether I wrote correctly the loop. And most importantly, I believe that this can be done but much shorter and more convenient. I would be grateful for any idea.

user_3073930 · Answer

You could also do this with a merge

merged = df.merge(df.drop(columns='end'), on='id', suffixes=('', '_y'))
live_projects = merged[merged.begin<=merged.begin_y]
result = live_projects.groupby(['id','begin','end']).count().reset_index()
result.rename(columns={'begin_y':'new'}, inplace=True)

Create new variable based on groupby & value comparison

1 Answers