Эффективный способ случайного выбора всех строк в кадре данных pandas, соответствующих значению столбца

Question

У меня есть кадр данных pandas, содержащий около 2 миллионов строк, который выглядит как в следующем примере

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

Я хочу создать три подмножества этих данных, чтобы идентификатор столбца был взаимоисключающим.И каждый из поднабора включает в себя все строки, соответствующие столбцу идентификатора в основном кадре данных.

На данный момент я случайно перетасовываю столбец идентификаторов и выбираю уникальные идентификаторы в качестве списка.Используя этот список, я выбираю все строки, которые с фрейма данных принадлежат части списка.

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

Я думаю, что замедление наступает во вложенном списке isin внутри среза loc.Я попробовал другой подход, используя numpy и логический индекс, который, кажется, удваивает скорость.

Сначала настройте фрейм данных.Я не был уверен, сколько у вас было уникальных предметов, поэтому я выбрал 50. Я также не был уверен, сколько столбцов так произвольно выбрал 10000 столбцов и строк.

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

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

Затем я пытаюсь использовать в основном пустые массивы и избегатьвложенный список с использованием логического индекса.

# 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, времена были 817 мс, код выше работает на 445 мс.

Не уверен, если этопомогает.Хороший вопрос, спасибо.