pandas - 选择符合条件的连续行对

时间:2016-12-14 04:52:45

标签: python pandas math

我有一个看起来像这样的数据框

>>> 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_0state_1。但这只返回第一个条目:

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

我应该如何修复逻辑,以便返回所有连续的行?在熊猫中有更好的方法吗?

2 个答案:

答案 0 :(得分:4)

我使用像这样的功能

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 

答案 1 :(得分:2)

这是一个适合我的解决方案 - 但仅适用于数字行索引。我让你的数据框架更有趣,现在它有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]