Summing row values after a groupby but based on a dictionary condition?

Question

I am trying to figure out how to add row entries of the numeric columns(supply,demand) . I am at a complete loss. My initial thoughts are to do this with a dictionary where i specify conditions within a dictionary(or perhaps there is another more efficient way to do this) ?. My second thought was to separate the data set such that each part has its own dataframe and then i use a bunch of conditionals but this would my last option since my dataset has about 10000 parts. In this example I only want to combine the supply and demand of rows with Planning location China and/or CH_China with plants US and/or China. Dataset:

df = pd.DataFrame(data={'Plant': ['China','New Zealand','US','US','Mexico'], 'Planning location': ['China','New Zealand','CH_China','US','China' ],'Part':[1,1,1,1,1],'Demand':[-5,6,7,8,9],'Supply':[4,10,-2,5,6]})

Expected result:

df_result = pd.DataFrame(data={'Plant': ['New Zealand','China & US','US','Mexico'], 'Planning location': ['New Zealand','China & CH_China','US','China' ],'Part':[1,1,1,1],'Demand':[6,2,8,9],'Supply':[10,2,5,6]})

Answer 1

I'm assuming you want to group US/China by Part number:

plant_is_china_us = df["Plant"].isin(["China", "US"])
planning_is_china = df["Planning location"].isin(["China", "CH_China"])

df1 = (
    df[plant_is_china_us & planning_is_china]
    .groupby("Part", as_index=False)
    .sum()
)
df1[["Plant", "Planning location"]] = ["China & US", "China & CH_China"]

df2 = df[~(plant_is_china_us & planning_is_china)]
print(pd.concat([df2, df1]).reset_index(drop=True).to_markdown())

Prints:

	Plant	Planning location	Part	Demand	Supply
0	New Zealand	New Zealand	1	6	10
1	US	US	1	8	5
2	Mexico	China	1	9	6
3	China & US	China & CH_China	1	2	2