过滤 Pandas DataFrame 中的行

Question

I am looking for a way to filter rows in a DataFrame. I have the following data:

data = [
    {'year':2015, 'v1':'str1', 'v2':'str2', 'v3':'str3', 'val': 6}, 
    {'year':2016, 'v1':'str1', 'v2':'str2', 'v3':'str3', 'val': 5}, 
    {'year':2017, 'v1':'str1', 'v2':'str2', 'v3':'str3', 'val': 3},
    {'year':2015, 'v1':'str11', 'v2':'str2', 'v3':'str3', 'val': 4},
    {'year':2016, 'v1':'str11', 'v2':'str2', 'v3':'str3', 'val': 9},
    {'year':2017, 'v1':'str12', 'v2':'str2', 'v3':'str3', 'val': 1},
    {'year':2016, 'v1':'str1', 'v2':'str21', 'v3':'str3', 'val': 9},
    {'year':2017, 'v1':'str1', 'v2':'str21', 'v3':'str3', 'val': 7},
    {'year':2018, 'v1':'str1', 'v2':'str21', 'v3':'str3', 'val': 8},
    {'year':2015, 'v1':'str1', 'v2':'str2', 'v3':'str31', 'val': 6}, 
    {'year':2016, 'v1':'str1', 'v2':'str2', 'v3':'str31', 'val': 5},
    {'year':2016, 'v1':'str1', 'v2':'str2', 'v3':'str31', 'val': 6}, 
    {'year':2017, 'v1':'str1', 'v2':'str2', 'v3':'str31', 'val': 3},
    {'year':2018, 'v1':'str1', 'v2':'str2', 'v3':'str31', 'val': 4}
]

The filtering rule: if there are not at least three subsequent years, starting with 2015, with rows which match in v1, v2 and v3, then those rows should be removed. The rows which match in v1, v2 and v3 for at least three subsequent years from 2015 on, should be kept.

The expected output after filtering for the example above is:

import pandas as pd
df = pd.DataFrame(data)
# filtering step
print(df)

    year     v1     v2     v3  val
0   2015   str1   str2   str3    6
1   2016   str1   str2   str3    5
2   2017   str1   str2   str3    3
3   2015   str1   str2  str31    6
4   2016   str1   str2  str31    5
5   2016   str1   str2  str31    6
6   2017   str1   str2  str31    3
7   2018   str1   str2  str31    4

Any ideas?

Answer 1

You can chain two groupby + filter

v = ['v1', 'v2', 'v3']

(df.groupby(v).filter(lambda s: 2015 in s['year'].values)
   .groupby(v).filter(lambda s: s.year.nunique() >= 3) and s.year.diff().isin([0, 1, np.nan]).all())

   year    v1    v2     v3  val
0  2015  str1  str2   str3    6
1  2016  str1  str2   str3    5
2  2017  str1  str2   str3    3
3  2015  str1  str2  str31    6
4  2016  str1  str2  str31    5
5  2016  str1  str2  str31    6
6  2017  str1  str2  str31    3
7  2018  str1  str2  str31    4

Answer 2

I feel like we can short the filter as below

df.groupby(['v1','v2','v3']).filter(lambda x : pd.Series([2015,2016,2017]).isin(x['year']).all())
Out[142]: 
    year    v1    v2     v3  val
0   2015  str1  str2   str3    6
1   2016  str1  str2   str3    5
2   2017  str1  str2   str3    3
9   2015  str1  str2  str31    6
10  2016  str1  str2  str31    5
11  2016  str1  str2  str31    6
12  2017  str1  str2  str31    3
13  2018  str1  str2  str31    4