Как посчитать количество элементов до определенного значения?

Дан df:

a = pd.DataFrame([{'id': 1, 'b': 1, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 4, 'd': 1},

                  {'id': 2, 'b': 1, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 4, 'd': 1},])

В каждом id будет только одна строка, где b == 1 и ниже 5 строк с b == 2 Нужно, сгруппировав по id, среди строк, где b == 2, посчитать количество строк, где d == 0 и записать в колонку k, но сделать это следующим образом:

  1. берем первую встретившуюся строку с b == 2, если у нее d == 0, то k = 1
  2. далее вторую строку с b == 2, если у нее d == 0, то k уже будет = 2
  3. и так все 5 строк с b == 2 пока не попадется d == 1
  4. то есть если в первой же строке с b == 2 d == 1, то k = 0 и следующие строки с b == 2 в этом id мы уже не рассматриваем.

В итоге должно получится вот так:

        pd.DataFrame([{'id': 1, 'b': 1, 'd': 1, 'k': 3},
                      {'id': 1, 'b': 2, 'd': 0, 'k': 3},
                      {'id': 1, 'b': 2, 'd': 0, 'k': 3},
                      {'id': 1, 'b': 2, 'd': 0, 'k': 3},
                      {'id': 1, 'b': 2, 'd': 1, 'k': 3},
                      {'id': 1, 'b': 2, 'd': 1, 'k': 3},
                      {'id': 1, 'b': 4, 'd': 1, 'k': 3},
    
                      {'id': 2, 'b': 1, 'd': 1, 'k': 1},
                      {'id': 2, 'b': 2, 'd': 0, 'k': 1},
                      {'id': 2, 'b': 2, 'd': 1, 'k': 1},
                      {'id': 2, 'b': 2, 'd': 0, 'k': 1},
                      {'id': 2, 'b': 2, 'd': 1, 'k': 1},
                      {'id': 2, 'b': 2, 'd': 1, 'k': 1},
                      {'id': 2, 'b': 4, 'd': 1, 'k': 1}])

Судя по ответу от @strawdog на мой предыдущий похожий вопрос Pandas посчитать количество элементов равное заданному значению, мы также должны использовать:

a["k"] = a.groupby("id").apply(lambda x: (x["b"]==2)&(x["d"] == 0)).groupby("id").transform(sum).to_list()

Но там необходимо было посчитать количество в целом, а здесь нужно посчитать количество до первого встретившегося d == 1, как можно добавить такое условие, не могу понять

P.S. Возможно, стоит просто сгруппировать по id и b == 2:

DD = a.query("b == 2").groupby('id')
AA = 0

Далее циклом перебирать каждую строку в каждом id, если d == 0, то обновить счетчик AA += 1, если d != 0, то break.

Но мягко говоря, такое решение кажется неадекватным, и что-то мне подсказывает, что можно решить эту ситуацию более удобными и быстрыми встроенными в Pandas методами.

P.S. @strawdog, если меняю входящий df на:

a = pd.DataFrame([{'id': 1, 'b': 1, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 4, 'd': 1},

                  {'id': 2, 'b': 1, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 4, 'd': 1}])

то на выходе получаю: введите сюда описание изображения

Но в столбце k должны быть 1 и 0 для 1 и 2 id соответственно:


P.S. Пробую делать вот таким корявым способом:

a = pd.DataFrame([{'id': 1, 'b': 1, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 4, 'd': 0},

                  {'id': 2, 'b': 1, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 4, 'd': 0},])

DD = a.query("b == 2").groupby('id')
for i in DD['d']:
    K = 0
    for index, j in enumerate(i):
        if index == 1:
    
            for m in j:
                if m == 0:
                    K += 1
                else:
                    break

    print(K)

  >>> 5
  >>> 1

То считается правильно вроде как, остается только эти цифры для каждого id добавить в колонку k. Но данный способ мне не нравится, он наверняка относительно затратный + я не до конца понимаю зачем мне приходится столько циклов перебирать)) С Pandas у меня проблемы, поэтому заранее сорри.


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

Автор решения: strawdog

Довольно сложное условие для векторизованной обработки фрейма. Я думаю, вам нужно как-то по-другому готовить данные перед помещением их в таблицу. Пока же я ничего более изящного, чем код ниже, придумать не смог:

import pandas as pd
import numpy as np

a = pd.DataFrame([{'id': 1, 'b': 1, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 0},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 2, 'd': 1},
                  {'id': 1, 'b': 4, 'd': 1},
                  {'id': 2, 'b': 1, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 0},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 2, 'd': 1},
                  {'id': 2, 'b': 4, 'd': 1},])

a["k"] =  a[((a["b"]==1)&(a["d"]==1))|((a["b"]==2)&(a["d"]==1))]["b"].astype("bool")
a["k"] = a["k"].fillna(0).cumsum()

res = pd.DataFrame()
for i, g in a.groupby("k"):
    t = g.copy()
    zeros = t["d"].to_list().count(0)
    t["k"] = t["d"].to_list().count(0)    
    res = pd.concat([res, t])
    
res["k"] = res["k"].replace(0, np.nan).fillna(method="ffill").astype("int")

res:

    id  b  d  k
0    1  1  1  3
1    1  2  0  3
2    1  2  0  3
3    1  2  0  3
4    1  2  1  3
5    1  2  1  3
6    1  4  1  3
7    2  1  1  1
8    2  2  0  1
9    2  2  1  1
10   2  2  0  1
11   2  2  1  1
12   2  2  1  1
13   2  4  1  1
→ Ссылка
Автор решения: MaxU

Если я правильно понял логику алгоритма:

a["k"] = (a["id"]
          .map(a.assign(x=a["d"].eq(0).astype(int))
                .query("b==2")
                .groupby("id")
                ["x"]
                .apply(lambda x: x.cumprod().cumsum().max())))

результат:

In [171]: a
Out[171]:
    id  b  d  k
0    1  1  1  3
1    1  2  0  3
2    1  2  0  3
3    1  2  0  3
4    1  2  1  3
5    1  2  1  3
6    1  4  1  3
7    2  1  1  1
8    2  2  0  1
9    2  2  1  1
10   2  2  0  1
11   2  2  1  1
12   2  2  1  1
13   2  4  1  1

пошагово:

сначала создаем вспомогательный виртуальный столбец x со значением 1 в тех строках, где d == 0 и 0 во всех остальных строках и фильтруем строки оставляя только те, где b == 2:

In [173]: a.assign(x=a["d"].eq(0).astype(int)).query("b==2")
Out[173]:
    id  b  d  k  x
1    1  2  0  3  1
2    1  2  0  3  1
3    1  2  0  3  1
4    1  2  1  3  0
5    1  2  1  3  0
8    2  2  0  1  1
9    2  2  1  1  0
10   2  2  0  1  1
11   2  2  1  1  0
12   2  2  1  1  0

теперь группируем записи по id и считаем кумулятивное произведение по столбцу x в каждой группе - это обнулит все значения в группе, начиная с первой встретившийся строки у которой d != 0:

In [174]: a.assign(x=a["d"].eq(0).astype(int)).query("b==2").groupby("id")["x"].apply(lambda x: x.cumprod())
Out[174]:
1     1
2     1
3     1
4     0
5     0
8     1
9     0
10    0
11    0
12    0
Name: x, dtype: int64

далее считаем кумулятивную сумму значения, получившегося на предыдущем шаге для каждой группы:

In [175]: a.assign(x=a["d"].eq(0).astype(int)).query("b==2").groupby("id")["x"].apply(lambda x: x.cumprod().cumsum())
Out[175]:
1     1
2     2
3     3
4     3
5     3
8     1
9     1
10    1
11    1
12    1
Name: x, dtype: int64

и берем максимальное значение для каждой группы:

In [176]: a.assign(x=a["d"].eq(0).astype(int)).query("b==2").groupby("id")["x"].apply(lambda x: x.cumprod().cumsum().max())
Out[176]:
id
1    3
2    1
Name: x, dtype: int64

в итоге используем получившийся ряд, подставляя вместо id соответствующее значение из этого ряда и присваивая результат новому столбцу (см. решение целиком выше).

→ Ссылка