Grouping by and applying lambda with condition for the first row - Pandas

user_11974225 0
Viewed 409

I have a data frame with IDs, and choices that have made by those IDs. The alternatives (choices) set is a list of integers: [10, 20, 30, 40]. Note: That's important to use this list. Let's call it 'choice_list'.

This is the data frame:

ID  Choice
1   10
1   30
1   10
2   40
2   40
2   40
3   20
3   40
3   10

I want to create a variable for each alternative: '10_Var', '20_Var', '30_Var', '40_Var'. At the first row of each ID, if the first choice was '10' for example, so the variable '10_Var' will get the value 0.6 (some parameter), and each of the other variables ('20_Var', '30_Var', '40_Var') will get the value (1 - 0.6) / 4. The number 4 stands for the number of alternatives.

Expected result:

ID  Choice  10_Var  20_Var  30_Var  40_Var
1   10      0.6     0.1     0.1     0.1
1   30              
1   10              
2   40      0.1     0.1     0.1     0.6
2   40              
2   40              
3   20      0.1     0.6     0.1     0.1
3   40              
3   10
pandas
3 Answers

you can use np.where to do this. It is efficient that df.where

df = pd.DataFrame([['1', 10], ['1', 30], ['1', 10], ['2', 40], ['2', 40], ['2', 40], ['3', 20], ['3', 40], ['3', 10]], columns=('ID', 'Choice'))

choices = np.unique(df.Choice)
for choice in choices:
    df[f"var_{choice}"] = np.where(df.Choice==choice, 0.6, (1 - 0.6) / 4)

df

Result

    ID  Choice  var_10  var_20  var_30  var_40
0  1      10     0.6     0.1     0.1     0.1
1  1      30     0.1     0.1     0.6     0.1
2  1      10     0.6     0.1     0.1     0.1
3  2      40     0.1     0.1     0.1     0.6
4  2      40     0.1     0.1     0.1     0.6
5  2      40     0.1     0.1     0.1     0.6
6  3      20     0.1     0.6     0.1     0.1
7  3      40     0.1     0.1     0.1     0.6
8  3      10     0.6     0.1     0.1     0.1

Edit To set values to 1st row of group only

df = pd.DataFrame([['1', 10], ['1', 30], ['1', 10], ['2', 40], ['2', 40], ['2', 40], ['3', 20], ['3', 40], ['3', 10]], columns=('ID', 'Choice'))
df=df.set_index("ID")

## create unique index for each row if not already
df = df.reset_index()

choices = np.unique(df.Choice)

## get unique id of 1st row of each group
grouped = df.loc[df.reset_index().groupby("ID")["index"].first()]

## set value for each new variable 
for choice in choices:
    grouped[f"var_{choice}"] = np.where(grouped.Choice==choice, 0.6, (1 - 0.6) / 4)

pd.concat([df, grouped.iloc[:, -len(choices):]], axis=1)
user_28516850

We can use insert o create the rows based on the unique ID values ​​obtained through Series.unique.We can also create a mask to fill only the first row using np.where. At the beginning sort_values ​​is used to sort the values ​​based on the ID. You can skip this step if your data frame is already well sorted (like the one shown in the example):

df=df.sort_values('ID')
n=df['Choice'].nunique()
mask=df['ID'].ne(df['ID'].shift())
for choice in df['Choice'].sort_values(ascending=False).unique():
    df.insert(2,column=f'{choice}_Var',value=np.nan)
    df.loc[mask,f'{choice}_Var']=np.where(df.loc[mask,'Choice'].eq(choice),0.6,0.4/n)
print(df)

       ID  Choice  10_Var  20_Var  30_Var  40_Var
0   1      10     0.6     0.1     0.1     0.1
1   1      30     NaN     NaN     NaN     NaN
2   1      10     NaN     NaN     NaN     NaN
3   2      40     0.1     0.1     0.1     0.6
4   2      40     NaN     NaN     NaN     NaN
5   2      40     NaN     NaN     NaN     NaN
6   3      20     0.1     0.6     0.1     0.1
7   3      40     NaN     NaN     NaN     NaN
8   3      10     NaN     NaN     NaN     NaN
user_118842370

A mix of numpy and pandas solution:

rows = np.unique(df.ID.values, return_index=1)[1]
df1 = df.loc[rows].assign(val=0.6)
df2 = (pd.crosstab([df1.index, df1.ID, df1.Choice], df1.Choice, df1.val, aggfunc='first')
         .reindex(choice_list, axis=1)
         .fillna((1-0.6)/len(choice_list)).reset_index(level=[1,2], drop=True))

pd.concat([df, df2], axis=1)

Out[217]:
   ID  Choice   10   20   30   40
0   1      10  0.6  0.1  0.1  0.1
1   1      30  NaN  NaN  NaN  NaN
2   1      10  NaN  NaN  NaN  NaN
3   2      40  0.1  0.1  0.1  0.6
4   2      40  NaN  NaN  NaN  NaN
5   2      40  NaN  NaN  NaN  NaN
6   3      20  0.1  0.6  0.1  0.1
7   3      40  NaN  NaN  NaN  NaN
8   3      10  NaN  NaN  NaN  NaN
user_101892140
Related