构建一个将数据框划分为组的函数

Question

我有兴趣创建一个执行以下操作的函数：

1. 接受2个参数：一个DataFrame和一个整数。
2. 在DF中添加一个名为“group”的列
3. 根据整数位置给每一行一个整数。组的数量应该是赋予函数的整数。
4. 如果行数不能通过给定的整数分割，则应该在组之间尽可能均匀地分割重新行。这是我遇到问题的部分。

这是我用来澄清我的意图的一个例子：

我想从这个DF中得到：

d = {'value': [1,2,3,4,5,6,7,8,9,10,11,12,13],}
df_init = pd.DataFrame(data=d)

通过这个功能：

wanted function(df_init,5)

对于这个精细的DF：

s = {'value': [1,2,3,4,5,6,7,8,9,10,11,12,13],'group':[1,1,1,2,2,2,3,3,3,4,4,5,5]}
df_finel = pd.DataFrame(data=d)

如果我能让问题更清楚，请告诉我如何解决问题。

Answer 1

使用np.array_split

In [5481]: [i for i, x in enumerate(np.array_split(np.arange(len(df)), 5), 1) for _ in x]
Out[5481]: [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 5, 5]

分配

In [5487]: df['group'] = [i for i, x in 
                          enumerate(np.array_split(np.arange(len(df)), 5), 1) for _ in x]

In [5488]: df
Out[5488]:
    value  group
0       1      1
1       2      1
2       3      1
3       4      2
4       5      2
5       6      2
6       7      3
7       8      3
8       9      3
9      10      4
10     11      4
11     12      5
12     13      5

详细

_原始df

In [5491]: df
Out[5491]:
    value
0       1
1       2
2       3
3       4
4       5
5       6
6       7
7       8
8       9
9      10
10     11
11     12
12     13

行为

In [5492]: np.array_split(np.arange(len(df)), 5)
Out[5492]:
[array([0, 1, 2]),
 array([3, 4, 5]),
 array([6, 7, 8]),
 array([ 9, 10]),
 array([11, 12])]