Python DataFrame单元格块移位优化：现有方案低效，求大数据集适配方案

Problem Description #

When working with a Pandas DataFrame, I encountered a column shifting formatting issue. Here's the example:

Original DataFrame #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': ['foo', 'bar', 'bar', 'foo', 'foo'],
    'B': ['spam', 'egg', 'egg', 'spam', 'spam'],
    'C': ['egg', '1', '2', 'egg', 'egg'],
    'D': [0, 0.84, 0.91, 3, 4],
    'E': [0, np.nan, np.nan, 0.14, -0.76]
})

Target DataFrame #

df = pd.DataFrame({
    'A': ['foo', 'bar', 'bar', 'foo', 'foo'],
    'B': ['spam', 'spam', 'spam', 'spam', 'spam'],
    'C': ['egg', 'egg', 'egg', 'egg', 'egg'],
    'D': [0, 1, 2, 3, 4],
    'E': [0.00, 0.84, 0.91, 0.14, -0.76]
})

My current fix code works but is not elegant and slow for large datasets:

columnShift = df.columns[-3:]
indexes = [i for i,x in enumerate(df.B) if x =='egg']
for i in indexes:
    for c in np.arange(1,len(columnShift)):
        df.loc[i,columnShift[-c]] = df.loc[i,columnShift[-(c+1)]]

I'm asking for a more optimized solution suitable for large datasets.

Answer #

Great question — looping through rows and columns with loc is definitely going to be slow on large datasets because Pandas isn't optimized for row-wise operations like that. Here are two much faster, vectorized approaches to solve this problem:

1. Boolean Mask + Vectorized Shift #

Since we only need to shift values when df['B'] == 'egg', we can create a mask for these rows and apply a vectorized shift to the relevant columns. This runs in C-level code, avoiding slow Python loops entirely:

import pandas as pd
import numpy as np

# Define the columns that need shifting (last 3 columns: C, D, E)
shift_cols = df.columns[-3:]

# Create a mask for rows where B equals 'egg'
mask = df['B'] == 'egg'

# Shift values right by 1 position for the masked rows
df.loc[mask, shift_cols] = df.loc[mask, shift_cols].shift(axis=1)

# Fill the now-empty 'B' column in masked rows with 'spam'
df.loc[mask, 'B'] = 'spam'

# Convert D column to numeric (since shifted rows had string values here originally)
df['D'] = pd.to_numeric(df['D'])

# Format E to 2 decimal places to match the target output
df['E'] = df['E'].round(2)

2. Conditional Assignment with where #

Another clean vectorized option uses pd.DataFrame.where to keep original values where the mask isn't met, and assign shifted values where it is:

shift_cols = df.columns[-3:]
mask = df['B'] == 'egg'

# Generate shifted values for the target columns
shifted_vals = df[shift_cols].shift(axis=1)

# Update the DataFrame: keep original values unless mask is true
df[shift_cols] = df[shift_cols].where(~mask, shifted_vals)
df['B'] = df['B'].where(~mask, 'spam')

# Clean up data types and formatting
df['D'] = pd.to_numeric(df['D'])
df['E'] = df['E'].fillna(0).round(2)

Why These Approaches Are Better #

• Speed: Vectorized operations avoid Python-level loops, which are the main bottleneck in your original code. For a DataFrame with 1 million rows, this will be 100x+ faster.
• Readability: The code clearly expresses the core logic (shift values where B is 'egg', fill in 'spam') without nested loops.
• Maintainability: It's easy to adjust if you need to change the number of columns to shift or the mask condition later.

Quick Notes #

• Don't forget to handle data type conversions (like converting D from string to numeric) since the shifted rows originally had string values in columns C and D.
• Adjust fillna() as needed if you have missing values outside the shifted rows in your actual dataset.

内容的提问来源于stack exchange，提问作者Alex