Блог

У меня есть фрейм данных pandas следующим образом:

 id variable year
1     a     2020
1     a     2030
1     a     2040
1     a     2050
1     b     2020
1     b     2030
1     b     2040
1     b     2050
1     c     2020
1     c     2030
1     c     2040
1     c     2050
 

Теперь для переменных в списке var = ['a','b'] я хочу сохранить только строки 2020 и 2030. В то время как для остальных переменных (например, c в приведенном выше примере) я хочу сохранить строки 2020, 2030 и 2040. Поэтому конечный фрейм данных должен быть:

 id variable year
1     a     2020
1     a     2030
1     b     2020
1     b     2030
1     c     2020
1     c     2030
1     c     2040
 

Как я могу этого добиться?

Используйте Series.isin с chain by | для побитового OR , amp; для побитового AND , а также ~ для инвертирования маски:

 var =  ['a','b']
var1 = [2020,2030]
var2 = [2020,2030,2040]

m1 = df['variable'].isin(var)
m2 = df['year'].isin(var1)
m3 = df['year'].isin(var2)

df = df[(m1 amp; m2) | (~m1 amp; m3)]
print (df)
    id variable  year
0    1        a  2020
1    1        a  2030
4    1        b  2020
5    1        b  2030
8    1        c  2020
9    1        c  2030
10   1        c  2040
 

Альтернативой предоставленному превосходному решению было бы использование метода запроса:

 search_criteria = "variable in ('a','b') and year in (2020, 2030) or variable=='c' and year in (2020, 2030, 2040)"
df.query(search_criteria)   

    id  variable    year
0   1      a       2020
1   1      a       2030
4   1      b       2020
5   1      b       2030
8   1      c       2020
9   1      c       2030
10  1      c       2040