从列值对应的数据帧中随机选择所有行的有效方法

Question

我有一个包含大约200万行的pandas数据帧，它看起来像下面的示例

ID  V1    V2   V3   V4    V5
12  0.2   0.3  0.5  0.03  0.9
12  0.5   0.4  0.6  0.7   1.8
01  3.8   2.9  1.1  1.6   1.5
17  0.9   1.2  1.8  2.6   9.0
02  0.2   0.3  0.5  0.03  0.9
12  0.5   0.4  0.6  0.7   1.8
07  3.8   2.9  1.1  1.6   1.5
19  0.9   1.2  1.8  2.6   9.0
19  0.5   0.4  0.6  0.7   1.8
06  3.8   2.9  1.1  1.6   1.5
17  0.9   1.2  1.8  2.6   9.0
18  0.9   1.2  1.8  2.6   9.0

我想创建此数据的三个子集，以便列ID互斥。并且每个子集包括与主数据帧中的ID列相对应的所有行。你知道吗

到现在为止，我随机地洗牌ID列并选择unique ID作为一个列表。使用这个列表，我从dataframe中选择属于列表一部分的ID的所有行。你知道吗

import numpy as np
import random 
distinct = list(set(df.ID.values))
random.shuffle(distinct)
X1, X2 = distinct[:1000000], distinct[1000000:2000000] 

df_X1 = df.loc[df['ID'].isin(list(X1))]

df_X2 = df.loc[df['ID'].isin(list(X2))]

对于较小的数据，这是可以预期的，但是对于较大的数据，运行甚至没有完成很多小时。有没有更有效的方法？感谢回应。你知道吗

Answer 1

我认为在loc切片中嵌套的isin列表中会出现减速。我尝试了一种不同的方法，使用numpy和一个看起来速度加倍的布尔索引。你知道吗

首先设置数据帧。我不知道你有多少独特的项目，所以我选择了50。我也不确定有多少列如此随意地选择了10000列和行。你知道吗

df = pd.DataFrame(np.random.randn(10000, 10000))

ID = np.random.randint(0,50,10000)
df['ID'] = ID

然后我尝试使用numpy数组，避免使用布尔索引的嵌套列表。你知道吗

# Create a numpy array from the ID columns
a_ID = np.array(df['ID'])

# use the numpy unique method to get a unique array
# a = np.unique(np.array(df['ID']))
a = np.unique(a_ID)

# shuffle the unique array
np.random.seed(100)
np.random.shuffle(a)

# cut the shuffled array in half
X1 = a[0:25]

# create a boolean mask
mask = np.isin(a_ID, X1)

# set the index to the mask
df.index = mask
df.loc[True]

当我在我的示例df上运行你的代码时，时间是817ms，上面的代码是445ms

不确定这是否有用。好问题，谢谢。你知道吗