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

【问题标题】：Split pandas DataFrame into approximately the same chunks将 pandas DataFrame 拆分为大致相同的块
【发布时间】：2021-04-23 14:52:55
【问题描述】：

我想将此 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 将出现在 2 个块中，因此这种分成 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 对象拆分为块时，一些解决方案就足够了。

【问题讨论】：

你能添加你理想的输出吗？它应该是由用户拆分的数据框列表吗？
当 n=2 时会发生什么？
@sophocles，感谢您的评论！它应该是一个大小大致相同的 DataFrame 列表，但一个用户应该只在一个结果 DataFrame 中。
@anky，这个很灵活，两种情况 [[0,1,2], [3, 4, 5]] 和 [[0,1], [2, 3, 4 , 5]] 就足够了。但当然有些用户会有 2 个价值和数百个。
您需要保留订单吗？例如您可以使用groupby("user") 分组，然后将组合并在一起以达到所需的最终块数。这样您就可以结束将原始数据框中不接近的用户分组在一起，这样可以吗？

标签： pandas dataframe

【解决方案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

【讨论】：

【解决方案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

【讨论】：

或者这个：df_list = {user:df for user, df in df_grouped}
@sophocles 感谢您的回复，但我在此解决方案中看不到很多块。我需要一个大小大致相同的 DataFrame 列表。每个 DataFrame 可能有多个用户，但单个用户的值不应该在多个 DataFrame 中。
这仅回答了部分问题，您需要按用户拆分数据帧并将它们重新组合成大致统一大小的块，至少据我所知
抱歉，我现在无法处理这些更改。我将不得不将更改推迟到晚些时候或明天。谢谢大家

【解决方案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]

【讨论】：