pandas - 将 pandas DataFrame 拆分成大致相同的 block

Question

我想将此 DataFrame 拆分为预定义数量的大小大致相同的 block :

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 很大(数百万行)，因此代码的效率应该或多或少。问题是某些用户应该只出现在其中一个 block 中。

例如如果 block 的数量是 3 那么:

• 第一个 block 应该有行 [0, 1]
• 第二个 block 应该有第 2 行，没有第 3 行，因为第 3 行是给用户 C 的
• 第三 block 应该有行 [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

这种分成 3 block 的分 block 是不正确的，因为用户 C 将出现在 2 个 block 中:

# 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 对象拆分成 block 时，一些解决方案应该就足够了。

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重新绑定(bind)

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