Найти длину полосы в пандах

Question

У меня есть фрейм данных pandas, где столбец отображает целочисленный временной индекс, и я хочу добавить столбец, в котором хранится информация о том, является ли строка частью полосы и какова длина полосы. Например, учитывая столбец time, я хотел бы вычислить столбец streak, например,

time    streak
0       3
1       3
2       3
4       2
5       2
5       2
9       1
11      1
11      1

Первые три строки являются частью полосы из трех, так как индексы 0,1,2 являются смежными. Следующие три строки имеют полосу 2, так как индексы 4,5 также являются смежными; Индекс 5 повторяется, но это не должно учитываться при определении длины полосы. Наконец, последние три строки не соприкасаются с чем-либо еще, поэтому они имеют полосу 1. Обратите внимание, что иногда несколько строк могут иметь одинаковые значения time. Мне нужно посчитать длину полосы в единицах времени, чтобы множественные записи не влияли на длину полосы, а строки с одинаковым индексом времени имели одинаковую длину полосы. Помните, что другие столбцы (не показаны) хранятся в кадре данных.

Как мне получить значение? Я пытался поиграться с groupby, shift и подобными функциями, но не очень далеко.

РЕДАКТИРОВАТЬ: извините, я забыл указать, что иногда индекс time может повторяться. Я расширил вопрос, чтобы учесть это.

Answer 1

Очень похоже на ответ Вена, просто используйте value_counts, что, на мой взгляд, немного более неприятно.

time = pd.Series([0, 1, 2, 4, 5, 9, 11])

# Give each row a streak id by incrementing whenever the difference isn't 1
streak = (time.diff() != 1).cumsum()

# Maps each id to the number of times the id occurs
result = streak.map(streak.value_counts())

print(result)
Out:
0    3
1    3
2    3
3    2
4    2
5    1
6    1

Изменить: вот решение для нового случая, добавленного к вопросу, где может быть повторное время. Обратите внимание, что мы используем diff > 1 для поиска новых полос; это использует предположение, что времена (не обязательно строго сейчас) увеличиваются целыми числами. Возможное дублирование означает, что нам нужно drop_duplicates перед тем, как определить количество идентификаторов полос для сопоставления.

time = pd.Series([0, 1, 2, 4, 5, 5, 9, 11, 11])

result = (time.diff() > 1).cumsum().map(
    (time.drop_duplicates().diff() > 1).cumsum().value_counts()
)

print(result)
Out:
0    3
1    3
2    3
3    2
4    2
5    2
6    1
7    1
8    1

Answer 2

Используя diff, найдите продолжение или нет (равное 1), затем cumsum с условием соответствия, затем мы используем groupby + transform szie

s=df.time.diff().fillna(1).ne(1).cumsum()
s.groupby(s).transform('size')
Out[396]: 
0    3
1    3
2    3
3    2
4    2
5    1
6    1
Name: time, dtype: int32