需要在pandas.explode（）之后索引子列表

Question

我需要跟踪在pd.explode()之后创建的项目的列表位置。虽然我只需要几行代码就可以做到这一点，比如在iterrows()中，但我正在尝试找到一种快速的方法来做到这一点

爆炸前

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

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

爆炸后

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

所需状态包括带有原始列表索引的idx列。如何正确地创建它

Answer 1

对于重复的索引值，首先从explode中删除ignore_index=True，以便可以通过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

另一个想法是通过列表长度展平range：

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