Python: разбиение DataFrame на основе числовой последовательности

Question

Я ищу реализацию Pythonic разделения DataFrame для панд на основе нескольких предварительно определенных числовых последовательностей в одном столбце (в данном примере state).

Пример:

sequence_1 = [4, 1, 5, 2]
sequence_2 = [3, 0]
test_data = pd.DataFrame({'state': [4, 1, 5, 2, 4, 1, 5, 2, 3, 0, 4, 1, 5, 2, 3, 0],
                      'output': [1, 1, 0, 1, 1, 3, 1, 1, 3, 2, 2, 2, 2, 0, 0, 0]})

Требуемый вывод: Разделить на

0 4 1
1 1 1
2 5 0
3 2 1

4 4 1
5 1 3
6 5 1
7 2 1

8 3 3
9 0 2

и т. Д.

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

Любая помощь будет принята с благодарностью!

Answer 1

Быстрый способ, если ваши данные в state упорядочены, как в вашем примере, состоит в том, чтобы перехватывать только первый элемент обеих последовательностей, а затем cumsum в groupby, например:

for name_g, df_g in test_data.groupby(((test_data.state == sequence_1[0])| 
                                       (test_data.state == sequence_2[0]) ).cumsum()):
    print (df_g)

Еще один общий способ - использовать shift, чтобы проверить, находится ли последовательность в правильном порядке, а затем получить кадры данных в list, например:

ser_seq1 = np.array([test_data.state.shift(-i) == val 
                     for i, val in enumerate(sequence_1)]).all(0)
list_df_seq1 = [test_data.loc[i:i+len(sequence_1)] 
                for i in test_data.index[ser_seq1]]

и то же самое с sequence_2