Как указать несколько условий внутри loc?

Question

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

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

Нужно, сгруппировав строки по колонке id, применить условие: если после строки в группе со страной = россия и месяцем = март следует строка со страной отличной от россии, то все строки этой группы выводим в результирующем фрейме, а если условие не выполняется, то строки этой группы не выводим.

Пытаюсь это решить так (но результат неверный):

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

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

Answer 1

Вариант решения для случаю когда нужно проверять на вхождения строки, удовлетворяющей условиям, которая встречается в группе после вхождения строки с "со страной = россия и месяцем = март":

def f(df):
    r = df["страна"].eq("россия") & df["месяц"].eq("март")
    if not r.any():
        return False
    return df.loc[r.idxmax() + 1:, "страна"].fillna("россия").ne("россия").any()

res = df.groupby("id").filter(f)

результат:

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

PS судя по описанию задачи в вопросе, группа 011 тоже должна присутствовать в результате.

Answer 2

Вариант решения для случая когда нужно проверять на вхождения строки, удовлетворяющей условиям, которая встречается в группе сразу и непосредственно за строкой "со страной = россия и месяцем = март":

def f(df):
    r = df["страна"].eq("россия") & df["месяц"].eq("март")
    if not r.any():
        return False
    x = (df
         .assign(c=df["страна"].shift(),   # предыдущее значение `страна`
                 m=df["месяц"].shift())    # предыдущее значение `месяц`
         .query("c == 'россия' and m == 'март' and страна != 'россия'"))
    return not x.empty

res = df.groupby("id").filter(f)

результат:

In [71]: res
Out[71]:
     фрукт   страна  id   месяц
0    груша  франция  01  январь
1   огурец   россия  01    март
2    вишня      сша  01  январь
5    арбуз   россия  03    март
6    груша      сша  03    март
11    дыня   россия   6    март
12   вишня  албания   6  январь
13   банан  андорра   6    июль

БЛОГ НА HUSL

Как указать несколько условий внутри loc?

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