有没有更快的方法将熊猫数据帧分成两个互补的部分？

Question

Good evening all,

I have a situation where I need to split a dataframe into two complementary parts based on the value of one feature.

What I mean by this is that for every row in dataframe 1, I need a complementary row in dataframe 2 that takes on the opposite value of that specific feature.

In my source dataframe, the feature I'm referring to is stored under column "773", and it can take on values of either 0.0 or 1.0 .

I came up with the following code that does this sufficiently, but it is remarkably slow. It takes about a minute to split 10,000 rows, even on my all-powerful EC2 instance.

data = chunk.iloc[:,1:776]
listy1 = []
listy2 = []

for i in range(0,len(data)):
    random_row = data.sample(n=1).iloc[0]
    listy1.append(random_row.tolist())

    if random_row["773"] == 0.0:
        x = data[data["773"] == 1.0].sample(n=1).iloc[0]
        listy2.append(x.tolist())

    else: 
        x = data[data["773"] == 0.0].sample(n=1).iloc[0]
        listy2.append(x.tolist())

df1 = pd.DataFrame(listy1)
df2 = pd.DataFrame(listy2)

Note: I don't care about duplicate rows, because this data is being used to train a model that compares two objects to tell which one is "better."

Do you have some insight into why this is so slow, or any suggestions as to make this faster?

Answer 1

A key concept in efficient numpy / scipy / pandas coding is using library-shipped vectorized functions whenever possible. Try to process multiple rows at once instead of iterate explicitly over rows. ie avoid for loops and .iterrows() .

The implementation provided is a little subtle in terms of indexing, but the vectorization thinking should be straightforward as follows:

1. Draw the main dataset at once.
2. The complementary dataset: draw the 0-rows at once, the complementary 1-rows at once, and then put them into the corresponding rows at once.

Code :

import pandas as pd
import numpy as np
from datetime import datetime

np.random.seed(52)  # reproducibility
n = 10000
df = pd.DataFrame(
    data={
        "773": [0,1]*int(n/2),
        "dummy1": list(range(n)),
        "dummy2": list(range(0, 10*n, 10))
    }
)

t0 = datetime.now()
print("Program begins...")
# 1. draw the main dataset
draw_idx = np.random.choice(n, n)  # repeatable draw
df_main = df.iloc[draw_idx, :].reset_index(drop=True)

# 2. draw the complementary dataset

# (1) count number of 1's and 0's
n_1 = np.count_nonzero(df["773"][draw_idx].values)
n_0 = n - n_1

# (2) split data for drawing
df_0 = df[df["773"] == 0].reset_index(drop=True)
df_1 = df[df["773"] == 1].reset_index(drop=True)

# (3) draw n_1 indexes in df_0 and n_0 indexes in df_1
idx_0 = np.random.choice(len(df_0), n_1)
idx_1 = np.random.choice(len(df_1), n_0)

# (4) broadcast the drawn rows into the complementary dataset
df_comp = df_main.copy()
mask_0 = (df_main["773"] == 0).values
df_comp.iloc[mask_0 ,:] = df_1.iloc[idx_1, :].values  # df_1 into mask_0
df_comp.iloc[~mask_0 ,:] = df_0.iloc[idx_0, :].values  # df_0 into ~mask_0

print(f"Program ends in {(datetime.now() - t0).total_seconds():.3f}s...")

Check

print(df_main.head(5))
   773  dummy1  dummy2
0    0      28     280
1    1      11     110
2    1      13     130
3    1      23     230
4    0      86     860

print(df_comp.head(5))
   773  dummy1  dummy2
0    1      19     190
1    0      74     740
2    0      28     280  <- this row is complementary to df_main
3    0      60     600
4    1      37     370

Efficiency gain : 14.23s -> 0.011s (ca. 128x)