CODEWITHSUNDEEP
pythonpython-3.xpandasdataframe
Adept
Adept

Reputation: 554

Pandas Dataframe keep rows where values of 2 columns are in a list of couples

I have a list of couples :

year_month = [(2020,8), (2021,1), (2021,6)]

and a dataframe df

| ID | Year | Month |
| 1  | 2020 | 1     |
|       ...         |
| 1  | 2020 | 12    |
| 1  | 2021 | 1     |
|       ...         |
| 1  | 2021 | 12    |
| 2  | 2020 | 1     |
|       ...         |
| 2  | 2020 | 12    |
| 2  | 2021 | 1     |
|       ...         |
| 2  | 2021 | 12    |
| 3  | 2021 | 1     |
|       ...         |

I want to select rows where Year and Month are corresponding to one of the couples in the year_month list :

Output df :

| ID | Year | Month |
| 1  | 2020 | 8     |
| 1  | 2021 | 1     |
| 1  | 2021 | 6     |
| 2  | 2020 | 8     |
| 2  | 2021 | 1     |
| 2  | 2021 | 6     |
| 3  | 2020 | 8     |
|       ...         |

Any idea on how to automate it, so I have only to change year_month couples ? I want to put many couples in year_month, so I want to keep a list of couples, and not to list all possibilities in df :

I don't want to do such :

df = df[((df['Year'] == 2020) & (df['Month'] == 8)) |
 ((df['Year'] == 2021) & (df['Month'] == 1)) | ((df['Year'] == 2021) & (df['Month'] == 6))]

Upvotes: 0

Views: 664

Answers (2)

sophocles
sophocles

Reputation: 13821

You can use a list comprehension and filter your dataframe with your list of tuples as below:

year_month = [(2020,8), (2021,1), (2021,6)]
df[[i in year_month for i in zip(df.Year,df.Month)]]

Which gives only the paired values back:

   ID  Year  Month
2   1  2021      1
6   2  2021      1
8   3  2021      1

Upvotes: 1

Chris
Chris

Reputation: 29742

One way using pandas.DataFrame.merge:

df.merge(pd.DataFrame(year_month, columns=["Year", "Month"]))

Output:

   ID  Year  Month
0   1  2021      1
1   2  2021      1
2   3  2021      1

Upvotes: 1

Related Questions