用相同的顺序在熊猫中重复几行

Question

我有这样的数据。

输入：

index colA   colB   colC 
0      A      1      QQQ  
1      A      1      WWW  
2      A      1      EEE  
3      A      1      RRR   
4      B      2      TTT   
5      B      2      YYY  
6      B      2      UUU  
7      C      1      III 

我想用相同的colA在colB时间重复值行。例如，行索引4-6的colA为B，在colB处的值为2，然后我想重复第4-6行两次( 4-6的顺序不变)，如输出所示。如果colB值为3，则重复三次，如果值为1，则不要重复。另外，我想添加一个列colC表示colA的索引，这样colA=A没有重复，那么num是0，colA=B有两个重复，num是1和2，.有没有一种优雅的方法可以在熊猫身上做到这一点？

输出：

index colA   colB   colC    num
0      A      1      QQQ     0
1      A      1      WWW     0
2      A      1      EEE     0
3      A      1      RRR     0
4      B      2      TTT     1
5      B      2      YYY     1
6      B      2      UUU     1
7      B      2      TTT     2
8      B      2      YYY     2
9      B      2      UUU     2
10     C      1      III     3

Answer 1

使用自定义函数按顺序重复(在许多组中慢)和keys参数，因此可以通过GroupBy.ngroup添加新列num

def f(x):
    pos = int(x.iat[0, x.columns.get_loc('colB')])
    return pd.concat([x] * pos, keys=range(pos))   

df = df.groupby('colA', group_keys=False, sort=False).apply(f).reset_index(level=0)

df['num'] = df.groupby(['level_0','colA'], sort=False).ngroup()
df = df.drop('level_0', axis=1)
print (df)
      colA  colB colC  num
index                     
0        A     1  QQQ    0
1        A     1  WWW    0
2        A     1  EEE    0
3        A     1  RRR    0
4        B     2  TTT    1
5        B     2  YYY    1
6        B     2  UUU    1
4        B     2  TTT    2
5        B     2  YYY    2
6        B     2  UUU    2
7        C     1  III    3

Answer 2

您可以使用index.repeat复制索引，使用GroupBy.cumcount计算"num“列：

(df
 .loc[df.index.repeat(df['colB'])]
 .assign(num=lambda d: d.groupby(['colA', 'colC']).cumcount().cumsum()) 
)

产出：

   index colA  colB colC  num
0      0    A     1  QQQ    0
1      1    A     1  WWW    0
2      2    A     1  EEE    0
3      3    A     1  RRR    0
4      4    B     2  TTT    0
4      4    B     2  TTT    1
5      5    B     2  YYY    1
5      5    B     2  YYY    2
6      6    B     2  UUU    2
6      6    B     2  UUU    3
7      7    C     1  III    3