31

I have a pandas.DataFrame with measurements taken at consecutive points in time. Along with each measurement the system under observation had a distinct state at each point in time. Hence, the DataFrame also contains a column with the state of the system at each measurement. State changes are much slower than the measurement interval. As a result, the column indicating the states might look like this (index: state):

1:  3
2:  3
3:  3
4:  3
5:  4
6:  4
7:  4
8:  4
9:  1
10: 1
11: 1
12: 1
13: 1

Is there an easy way to retrieve the indices of each segment of consecutively equal states. That means I would like to get something like this:

[[1,2,3,4], [5,6,7,8], [9,10,11,12,13]]

The result might also be in something different than plain lists.

The only solution I could think of so far is manually iterating over the rows, finding segment change points and reconstructing the indices from these change points, but I have the hope that there is an easier solution.

languitar
  • 6,554
  • 2
  • 37
  • 62

2 Answers2

50

One-liner:

df.reset_index().groupby('A')['index'].apply(np.array)

Code for example:

In [1]: import numpy as np

In [2]: from pandas import *

In [3]: df = DataFrame([3]*4+[4]*4+[1]*4, columns=['A'])
In [4]: df
Out[4]:
    A
0   3
1   3
2   3
3   3
4   4
5   4
6   4
7   4
8   1
9   1
10  1
11  1

In [5]: df.reset_index().groupby('A')['index'].apply(np.array)
Out[5]:
A
1    [8, 9, 10, 11]
3      [0, 1, 2, 3]
4      [4, 5, 6, 7]

You can also directly access the information from the groupby object:

In [1]: grp = df.groupby('A')

In [2]: grp.indices
Out[2]:
{1L: array([ 8,  9, 10, 11], dtype=int64),
 3L: array([0, 1, 2, 3], dtype=int64),
 4L: array([4, 5, 6, 7], dtype=int64)}

In [3]: grp.indices[3]
Out[3]: array([0, 1, 2, 3], dtype=int64)

To address the situation that DSM mentioned you could do something like:

In [1]: df['block'] = (df.A.shift(1) != df.A).astype(int).cumsum()

In [2]: df
Out[2]:
    A  block
0   3      1
1   3      1
2   3      1
3   3      1
4   4      2
5   4      2
6   4      2
7   4      2
8   1      3
9   1      3
10  1      3
11  1      3
12  3      4
13  3      4
14  3      4
15  3      4

Now groupby both columns and apply the lambda function:

In [77]: df.reset_index().groupby(['A','block'])['index'].apply(np.array)
Out[77]:
A  block
1  3          [8, 9, 10, 11]
3  1            [0, 1, 2, 3]
   4        [12, 13, 14, 15]
4  2            [4, 5, 6, 7]
jamis
  • 247
  • 3
  • 10
Zelazny7
  • 39,946
  • 18
  • 70
  • 84
  • 19
    This assumes that the values don't repeat in discontiguous segments -- for example, `DataFrame([3]*4+[4]*4+[1]*4 + [3]*4, columns=['A'])` will put the two groups of 3 into the same group. You could scan those for breaks, but that's just another version of the original problem. Maybe there's a way to get the pandas `groupby` to behave more like `itertools.groupby` here, though. – DSM Jan 16 '13 at 14:27
  • 1
    Thanks, your second solution works well. I actually have the situation described by DSM. – languitar Jan 17 '13 at 17:29
  • How might this be done if your were to want to group by some deviation (e.g. groups contain values where all values are within +-1 of adjacent values in the original set) – ryanjdillon May 20 '14 at 09:22
  • @ryanjdillon you can try `df['block'] = (df.A.diff(1).abs() > 1).cumsum()`. but that only check the difference of adjacent values, so if you would have `DataFrame([3]*4+[4]*4+[5]*1+[6]*3, columns=['A'])`, it will still produce just single block. – kekert Sep 17 '16 at 23:59
  • It's worth noting that you don't actually need the "astype(int)" there--pandas is perfectly happen to sum Boolean values. – MTKnife May 05 '21 at 00:20
  • This blog post explains one of the solutions given https://towardsdatascience.com/pandas-dataframe-group-by-consecutive-same-values-128913875dba – waithira Jan 09 '23 at 11:42
19

You could use np.diff() to test where a segment starts/ends and iterate over those results. Its a very simple solution, so probably not the most performent one.

a = np.array([3,3,3,3,3,4,4,4,4,4,1,1,1,1,4,4,12,12,12])

prev = 0
splits = np.append(np.where(np.diff(a) != 0)[0],len(a)+1)+1

for split in splits:
    print np.arange(1,a.size+1,1)[prev:split]
    prev = split

Results in:

[1 2 3 4 5]
[ 6  7  8  9 10]
[11 12 13 14]
[15 16]
[17 18 19]
Rutger Kassies
  • 61,630
  • 17
  • 112
  • 97