Is there a pythonic way to add an enumerating column while exploding a list column in pandas?

user_2153636 0
Viewed 123

Consider the following DataFrame:

>>> df = pd.DataFrame({'A': [1,2,3], 'B':['abc', 'def', 'ghi']}).apply({'A':int, 'B':list})
>>> df
   A          B
0  1  [a, b, c]
1  2  [d, e, f]
2  3  [g, h, I]

This is one way to get the desired result:

>>> df['B'] = df['B'].apply(enumerate).apply(list)
>>> df = df.explode('B', ignore_index=True)
>>> df['B'] = pd.Series(df['B'], index=['B1', 'B2'])})
>>> df.droplevel(0, axis=1)
   A  B1 B2
0  1   0  a
1  1   1  b
2  1   2  c
3  2   0  d
4  2   1  e
5  2   2  f
6  3   0  g
7  3   1  h
8  3   2  i

Is there a neater way?

pandas enumeration dataframe python
1 Answers

A groupby on the index is an option:

df.explode('B').assign(
  B1 = lambda df: df.groupby(level=0).cumcount())

   A  B  B1
0  1  a   0
0  1  b   1
0  1  c   2
1  2  d   0
1  2  e   1
1  2  f   2
2  3  g   0
2  3  h   1
2  3  i   2

you can always reset the index, if you have no use for it:

df.explode('B').assign(
  B1 = lambda df: df.groupby(level=0).cumcount()).reset_index(drop=True)

   A  B  B1
0  1  a   0
1  1  b   1
2  1  c   2
3  2  d   0
4  2  e   1
5  2  f   2
6  3  g   0
7  3  h   1
8  3  i   2

Since pandas version 1.3.0 you can use multiple columns with explode out of the box:

df.assign(
  B1 = df.B.apply(len).apply(range)).explode(['B', 'B1'], ignore_index = True))

   A  B B1
0  1  a  0
1  1  b  1
2  1  c  2
3  2  d  0
4  2  e  1
5  2  f  2
6  3  g  0
7  3  h  1
8  3  i  2

I think a faster option would be to run the reshaping outside Pandas, and then rejoin back to the dataframe (of course only tests can affirm/deny this):

from itertools import chain
# you can use np.concatenate instead
# np.concatenate(df.B)
flattened = chain.from_iterable(df.B)
index = df.index.repeat([*map(len, df.B)])
flattened = pd.Series(flattened, index, name = 'B1')
(pd.concat([df.A, flattened], axis=1)
  .assign(B2 = lambda df: df.groupby(level=0).cumcount())
)

   A B1  B2
0  1  a   0
0  1  b   1
0  1  c   2
1  2  d   0
1  2  e   1
1  2  f   2
2  3  g   0
2  3  h   1
2  3  i   2
user_71757130
Related