将pandas DataFrame拆分为大致相同的块

Question

我想将此DataFrame拆分为预定义数量的块，大小大致相同：

import pandas as pd
df = pd.DataFrame({
"user": ["A", "A", "B", "C", "C", "C"],
"value": [0.3, 0.4, 0.5, 0.6, 0.7, 0.8]})

#     user  value
# 0      A    0.3
# 1      A    0.4
# 2      B    0.5
# 3      C    0.6
# 4      C    0.7
# 5      C    0.8

DataFrame很大(数百万行(，因此代码的效率应该或多或少。问题是，某个用户应该只出现在其中一个块中。

例如，如果块的数量是3，那么：

• 第一个区块应该有行[0，1]
• 第二个区块应该有第2行，而不是第3行，因为第3行是为用户C准备的
• 第三块应该有行[3，4，5]

# Chunk #1 (DataFrame):
# 0      A    0.3
# 1      A    0.4
# Chunk #2 (DataFrame):
# 2      B    0.5
# Chunk #3 (DataFrame):
# 3      C    0.6
# 4      C    0.7
# 5      C    0.8

由于用户C将以两个块的形式出现，因此这种3块的分块是不正确的：

# Chunk #1 (DataFrame):
# 0      A    0.3
# 1      A    0.4
# Chunk #2 (DataFrame):
# 2      B    0.5
# 3      C    0.6
# Chunk #3 (DataFrame):
# 4      C    0.7
# 5      C    0.8

我认为，当我们先按用户执行groupby，然后将该DataFrameGroupBy对象分块时，一些解决方案就足够了。

Answer 1

您可以将我们的user列转换为分类列，并使用qcut进行统一高度装箱。不幸的是，qcut无法为不连续的分布找到唯一的bin边，因此如果一个用户被过度表示，您可能会遇到一些问题。你可以使用duplicates="drop"，但你不会总是有你要求的数字箱，因为有些会聚集在一起。

我想你必须写一些算法来进行适当的重新编译，但找不到现成的东西。

下面是pd.qcut的一个例子。

让我们构建一个伪数据集

user = np.random.choice(["A", "B", "C", "D", "E", "F", "G", "H"], 10000)
value = np.random.random(size=user.shape)
df = pd.DataFrame({"user": user, "value": value})
print(df.user.value_counts())
E    1329
C    1281
G    1277
F    1260
H    1231
D    1223
A    1205
B    1194
Name: user, dtype: int64

为每个唯一用户分配一个整数代码，并使用qcut重新绑定

codes = df.user.astype("category").cat.codes    
nbins = 3
df["bin"] = pd.qcut(codes, nbins, labels=False)
df.groupby("user").bin.value_counts()

让我们检查一下的结果

print(df.bin.value_counts())
1    3788
0    3629
2    2583
Name: bin, dtype: int64
print(df.groupby("user").bin.value_counts())
user  bin
A     0      1266
B     0      1158
C     0      1205
D     1      1255
E     1      1246
F     1      1287
G     2      1274
H     2      1309
Name: bin, dtype: int64

Answer 2

这就足够了吗？

df_grouped = df.groupby('user')
df_list = [df for user, df in df_grouped]

---

Out[1352]: 
[  user  value
0    A    0.3
1    A    0.4,
user  value
2    B    0.5,
user  value
3    C    0.6
4    C    0.7
5    C    0.8]

这在我的机器上运行得相对较快：

>>> df.shape
(7200000, 2)
>>> print(end - start)
0.532534122467041

Answer 3

您也可以尝试np.split，并在适当的条件下检查块是否与用户中的元素数量相同，然后在用户之间进行拆分，否则拆分时保留列表中的前n个用户：

def split_fun(data,n):
cond = len(set(data['user'])) == n
f = data['user'].factorize()[0]+1
if cond:
p = np.where(np.diff(f)>0)[0]+1
else:
p= np.where(np.diff((f>n).view('i1'))>0)[0]+1
return np.split(data,p)

---

样本运行：

split_fun(df,2)
[  user  value
0    A    0.3
1    A    0.4
2    B    0.5,
user  value
3    C    0.6
4    C    0.7
5    C    0.8]
split_fun(df,3)
[  user  value
0    A    0.3
1    A    0.4,
user  value
2    B    0.5,
user  value
3    C    0.6
4    C    0.7
5    C    0.8]