панды - Выбор пары последовательных строк, соответствующих критериям

Question

У меня есть dataframe, который выглядит, как этот

>>> a_df 
    state 
1 A 
2 B 
3 A 
4 B 
5 C 

Что я хотел бы сделать, это вернуть все последовательные строки соответствие определенной последовательности. Например, если эта последовательность равна ['A', 'B'], тогда должны быть возвращены строки, состояние которых A, за которым следует сразу B. В приведенном выше примере:

>>> cons_criteria(a_df, ['A', 'B']) 
    state 
1 A 
2 B 
3 A 
4 B 

Или, если выбранный массив ['A', 'B', 'C'], то выход должен быть

>>> cons_criteria(a_df, ['A', 'B', 'C']) 
    state 
3 A 
4 B 
5 C 

Я решил сделать это путем сохранения текущего состояния, а также следующее состояние:

>>> df2 = a_df.copy() 
>>> df2['state_0'] = a_df['state'] 
>>> df2['state_1'] = a_df['state'].shift(-1) 

Теперь я могу соответствовать по state_0 и state_1. Но это возвращает только самый первый вход:

>>> df2[(df2['state_0'] == 'A') & (df2['state_1'] == 'B')] 
    state 
1 A 
3 A 

Как исправить логику здесь, так что все последовательные строки возвращаются? Есть ли лучший способ приблизиться к этому в пандах?

Answer 1

Вот решение, которое работает для меня - но только для числовых индексов строк. Я сделал свой dataframe немного более интересным, теперь он имеет 2 ABC модели:

a_df=pd.DataFrame(['A','B','A','B','C','D','A','A','B','C','E'], 
        columns=["state"]) 

И вот картина матча:

pattern = ['A','B','C'] 

Это выражение находит набор исходных строк каждого шаблона:

starts = set(a_df[a_df['state']   =='A'].index) & 
     set(a_df[a_df['state'].shift(-1)=='B'].index) & 
     set(a_df[a_df['state'].shift(-2)=='C'].index) 
print(starts) 
# {2, 7} 

В общем:

starts = set.intersection(
      *[set(a_df[a_df['state'].shift(-i)==value].index) 
      for i,value in enumerate(pattern)]) 

Это выражение преобразует число строк в исходном диапазоны 3-строк и выбирает диапазоны строк:

result = [a_df.ix[range(i, i+3)] for i in starts] 
print(result) 
# [ state 
# 2  A 
# 3  B 
# 4  C, state 
# 7  A 
# 8  B 
# 9  C] 

В целом:

result = [a_df.ix[range(i, i+len(pattern))] for i in starts] 

Answer 2

Я хотел бы использовать функцию, как это

def match_slc(s, seq): 
    # get list, makes zip faster 
    l = s.values.tolist() 
    # count how many in sequence 
    k = len(seq) 
    # generate numpy array of rolling values 
    a = np.array(list(zip(*[l[i:] for i in range(k)]))) 
    # slice an array from 0 to length of a - 1 with 
    # the truth values of wether all 3 in a sequence match 
    p = np.arange(len(a))[(a == seq).all(1)] 
    # p tracks the beginning of a match, get all subsequent 
    # indices of the match as well. 
    slc = np.unique(np.hstack([p + i for i in range(k)])) 
    return s.iloc[slc] 

---

Демонстрация

s = pd.Series(list('ABABC')) 

print(match_slc(s, list('ABC')), '\n') 
print(match_slc(s, list('AB')), '\n') 

2 A 
3 B 
4 C 
dtype: object 

0 A 
1 B 
2 A 
3 B 
dtype: object