Need to index sub-lists after pandas.explode()

user_10396469 0
Viewed 164

I need to keep track of list positions of the items that are created after pd.explode(). While I can do this with few lines of code, like in iterrows(), I am trying to find fast Pandas way to do this.

before explode

df = pd.DataFrame({'foo':[['a','b','c'],['d'],['e', 'f']]})

         foo
0  [a, b, c]
1        [d]
2     [e, f]

after explode

df = df.explode('foo', ignore_index=True)
df['idx'] = [0,1,2,0,0,1]  # NEED TO REPLACE THIS LINE WITH SMART PANDAS FUNCTION

  foo  idx
0   a    0
1   b    1
2   c    2
3   d    0
4   e    0
5   f    1

The desired state includes idx column with indices of the original lists. How to create it properly?

pandas python
1 Answers

First remove ignore_index=True from explode for repeated index values, so possible use GroupBy.cumcount by index an last create default index:

df = df.explode('foo')

df['idx'] = df.groupby(level=0).cumcount()
#alternative
#df['idx'] = df.groupby(df.index).cumcount()
df = df.reset_index(drop=True)
print (df)
  foo  idx
0   a    0
1   b    1
2   c    2
3   d    0
4   e    0
5   f    1

Another idea is flatten ranges by lengths of lists:

L = [y for x in df['foo'] for y in range(len(x))]
df = df.explode('foo', ignore_index=True)
df['idx'] = L

print (df)
  foo  idx
0   a    0
1   b    1
2   c    2
3   d    0
4   e    0
5   f    1
user_29010020
Related