How to assign new variables within a group by - specifically np.where within pd.groupby

user_5024631 0
Viewed 34

I want to group my data by two different variables and then create a new column which state whether or nor "3" or "4" appears in the grouping, but I keep getting different types of errors. Here is my sample data:

d = {'id' : ["A","A","A","A","A","A","B","B","B","B","B","B"],
    'month' : [1,1,1,1,2,2,1,1,1,1,2,2],
    'week' : [1,2,3,4,1,2,1,2,3,4,1,2]}

example_df = pd.DataFrame(data = d)

I want to group by id and month and then within that grouping, create a new column titled has_3_or_4 when 3 or 4 appear in the column week

example_df = (example_df.assign(has_3_or_4 = example_df.groupby(['id', 'month'])
                .apply(lambda x: np.where(any(x.week.isin([3, 4])),'has_3_or_4', 'no_3_or_4'))))

But this returns this error:

TypeError: incompatible index of inserted column with frame index

I googled this and found a solution which says to include as_index = False within the groupby, which does indeed allow the code to run without any error, but the end result is not what I want, it just looks weird.

Does anyone have any idea how to do this? in R it is so straightforward - e.g. example_df %>% group_by(id, month) %>% mutate(has_3_or_4 = if_else(any(c(3,4) %in% week), 'has_3_or_4', 'no_3_or_4'))

pandas grouping dataframe group-by python
1 Answers

If I understand you correctly, you can create a new column with .transform:

example_df["has_3_or_4"] = example_df.groupby(["id", "month"])[
    "week"
].transform(lambda x: x.isin([3, 4]).any())

print(example_df)

Prints:

   id  month  week  has_3_or_4
0   A      1     1        True
1   A      1     2        True
2   A      1     3        True
3   A      1     4        True
4   A      2     1       False
5   A      2     2       False
6   B      1     1        True
7   B      1     2        True
8   B      1     3        True
9   B      1     4        True
10  B      2     1       False
11  B      2     2       False

EDIT: Using .assign:

example_df = example_df.assign(
    has_3_or_4=example_df.groupby(["id", "month"])["week"].transform(
        lambda x: np.where(any(x.isin([3, 4])), "has_3_or_4", "no_3_or_4")
    )
)
user_100359850
Related