Как отобрать группы строк по условию?

Question

Имеется фрейм данных:

import pandas as pd
import numpy as np
data = {'фрукт': ['груша','огурец','вишня', 'абрикос', 'груша','арбуз','груша','банан', 'груша', 'вишня','яблоко', 'груша', 'вишня', 'абрикос', 'груша', 'банан'],
'страна': ['франция','россия', 'сша','россия', 'франция','россия','сша', 'россия', 'гана','россия','россия', 'россия', 'албания', 'россия', np.nan, 'россия'],
'id': ['01','01','01','02','02','03','03','011', '011', '011','011', '6', '6', '5', '5', '5'],
'месяц': ['январь','март','январь','март', 'сентябрь','март','март', 'март', 'ноябрь', 'январь','январь', 'март', 'январь', 'март', 'март', 'апрель']        
}
dates = pd.DataFrame(data, columns = ['фрукт','страна', 'id', 'месяц'])
dates

Нужно оставить те записи групп id (то есть группируем по id и потом проверяем условие), в которых после строки со значением россия в колонке страна следует строка со значением страны отличным от значения россия и не NaN.

Как это можно сделать?

Пытаюсь это сделать так:

def pairs(ls):
    res = []
    for i in range(len(ls)-1):
        res.append([b[i],b[i+1]])
    return res

def orr(b,k):
    z = True
    for i in b:
        bulev = i == k
        z = z or bulev
    return z
df = dates.copy()
patterns = ["россия", "россия"]
res = df.groupby("id").filter(lambda x: orr(pairs(x["страна"].tolist()),patterns))

Но возникает ошибка:

'bool' object is not subscriptable

Ожидаемый результат:

Answer 1

UPDATE 2

res= pd.DataFrame()

for i, g in dates.groupby("id"):
    if any(g["страна"].str.contains("россия")) and g["страна"].nunique()>1:
        if any([x > min(g[g.loc[:,"страна"]=="россия"].index) for x in g[g.loc[:,"страна"]!="россия"].index]):
            res = res.append(g)

res:

      фрукт   страна   id     месяц
0     груша  франция   01    январь
1    огурец   россия   01      март
2     вишня      сша   01    январь
7     банан   россия  011      март
8     груша     гана  011    ноябрь
9     вишня   россия  011    январь
10   яблоко   россия  011    январь
3   абрикос   россия   02      март
4     груша  франция   02  сентябрь
5     арбуз   россия   03      март
6     груша      сша   03      март
11    груша   россия    6      март
12    вишня  албания    6    январь

Answer 2

res = (df
       .groupby("id")
       .filter(
           lambda x: x.iloc[x["страна"].eq("россия").argmax() + 1:, 
                            df.columns.get_loc("страна")]
                      .fillna("россия")
                      .ne("россия")
                      .any()))

результат:

In [97]: res
Out[97]:
      фрукт   страна   id     месяц
0     груша  франция   01    январь
1    огурец   россия   01      март
2     вишня      сша   01    январь
3   абрикос   россия   02      март
4     груша  франция   02  сентябрь
5     арбуз   россия   03      март
6     груша      сша   03      март
7     банан   россия  011      март
8     груша     гана  011    ноябрь
9     вишня   россия  011    январь
10   яблоко   россия  011    январь
11    груша   россия    6      март
12    вишня  албания    6    январь

Пояснения:

сначала находим порядковый индекс первого вхождения "россия" в группе - x["страна"].eq("россия").argmax()
выбираем все значения в столбце "страна" в группе начиная с найденного в предыдущем шаге индекса + 1. Т.е. выбираем строки начиная со следующей после первого вхождения "россия" в группе - x.loc[x["страна"].eq("россия").argmax() + 1:, df.columns.get_loc("страна")]
меняем NaN -> "россия": .fillna("россия")
проверяем есть ли в выбранных в предыдущих шагах значениях хотя бы одно значение отличное от "россия" - .ne("россия").any()

В том случае если значения индекса - монотонно возрастающая последовательность можно воспользоваться .loc[] и Series.idxmax():

res = (df
       .groupby("id")
       .filter(
           lambda x: x.loc[x["страна"].eq("россия").idxmax() + 1:, "страна"]
                      .fillna("россия")
                      .ne("россия")
                      .any()))

БЛОГ НА HUSL

Как отобрать группы строк по условию?

Ответы (2 шт):