我有一个带有ID列的数据框,我想在数据框中添加一列是每个唯一ID的索引。通过下面的示例,我可以使用2个for循环来做到这一点,即从ID计数中列出一个列表,将其转换为数组,然后将其添加到数据帧中。
但是,我希望改为使用pandas groupby()来简化此过程,并避免使用这两个for循环,但这不太有效。我试着使用groupby()和set_index(),但是我什么都无法工作。有人知道我如何使用pandas groupby()简化所有这些步骤吗?
df = pd.DataFrame({'ID': (4,4,4,5,6,7,7,7,7),
'Record': (1,4,5,1,1,1,5,7,8),
'color': ('b','r','r','r','r','b','b','b','b')})
ID Record color
4 1 b
4 4 r
4 5 r
5 1 r
6 1 r
7 1 b
7 5 b
7 7 b
7 8 b
##############################################
### Steps I did to add the make sub index ID to df
def convert_to_sub_index_array(count_list):
sub_ID_list = []
for c in count_list:
for i in range(c):
sub_ID_list.append(i+1)
return np.asarray(sub_ID_list)
ID_group=df.groupby(['ID']).count()
ID_count = ID_group['Record'].values.tolist()
sub_ID_index = convert_to_sub_index_array(ID_count)
df['ID_index'] = sub_ID_index
#################################################
ID Record color ID_index
4 1 b 1
4 4 r 2
4 5 r 3
5 1 r 1
6 1 r 1
7 1 b 1
7 5 b 2
7 7 b 3
7 8 b 4