优化含多条件判断的Pandas DataFrame循环计算逻辑

Hey there! I totally get your frustration—looping through 320k rows twice is going to crawl, especially with Python's native loops. Let's fix this by leveraging vectorized operations and spatial indexing, which will speed things up drastically (we're talking from seconds per iteration to milliseconds).

Why Your Original Code Is Slow #

Your current approach uses nested loops with df.at[] access, which has two big issues:

• O(N*M) time complexity: For ~5000 winner rows (320k / 64) and 320k total rows, that's 1.6e11 operations—way too much for Python to handle quickly.
• Slow per-operation overhead: df.at[] is designed for single-cell access, not bulk operations, so each call adds unnecessary slowdown.

Solution 1: Vectorized Operations with NumPy #

NumPy's vectorized calculations run in optimized C code, which avoids Python loop overhead entirely. Here's how to rewrite your function:

import pandas as pd
import numpy as np

# Load your data (only keep needed columns to save memory)
allDF = pd.read_csv('players-with-rank-data.csv', usecols=['x', 'y', 'rank'])
winnerDF = allDF[allDF['rank'] == 1].copy()  # Avoid SettingWithCopyWarning
radius_sq = 20000 ** 2  # Precompute squared radius (skip expensive sqrt!)

# Preprocess arrays once (outside loops for efficiency)
all_coords = allDF[['x', 'y']].to_numpy()
is_winner = allDF['rank'] == 1  # Boolean array for fast filtering

win_percentages = []
for _, row in winnerDF.iterrows():
    # Get center coordinates as a NumPy array
    center = np.array([row['x'], row['y']])
    # Calculate squared distance from center to ALL points (vectorized)
    dist_sq = np.sum((all_coords - center) ** 2, axis=1)
    # Filter points within the radius
    in_circle = dist_sq <= radius_sq
    # Count total and winner points in the circle
    total_in_circle = np.sum(in_circle)
    winners_in_circle = np.sum(is_winner[in_circle])
    # Calculate percentage (handle division by zero edge case)
    win_pct = winners_in_circle / total_in_circle if total_in_circle > 0 else 0.0
    win_percentages.append(win_pct)

# Add results to your winner dataframe (optional)
winnerDF['win_percentage'] = win_percentages

# Print or save results (batch IO is faster than per-iteration)
for pct in win_percentages:
    print(pct)
# winnerDF.to_csv('winner_results.csv', index=False)

Key Optimizations Here: #

• Precompute arrays: We convert coordinates and rank flags to NumPy arrays once, not inside the loop.
• Vectorized distance calculation: (all_coords - center) ** 2 uses broadcasting to compute distances for all points at once.
• Avoid square roots: Comparing squared distances to squared radius skips an expensive math operation.

Solution 2: KD-Tree Spatial Index (Even Faster!) #

For very large datasets, a KD-Tree will reduce the number of distance calculations by indexing spatial data. This is perfect if your radius is small compared to the overall data spread. We'll use scipy.spatial.KDTree:

import pandas as pd
import numpy as np
from scipy.spatial import KDTree

# Load data
allDF = pd.read_csv('players-with-rank-data.csv', usecols=['x', 'y', 'rank'])
winnerDF = allDF[allDF['rank'] == 1].copy()
radius = 20000

# Build KD-Tree from all coordinates (spatial index for fast neighbor searches)
tree = KDTree(allDF[['x', 'y']].to_numpy())
# Extract winner coordinates and rank flags
winner_coords = winnerDF[['x', 'y']].to_numpy()
is_winner = allDF['rank'].values == 1

# Batch query all winner points for neighbors within radius
# Returns a list of indices for each winner point
nearby_indices = tree.query_ball_point(winner_coords, r=radius)

win_percentages = []
for indices in nearby_indices:
    total_in_circle = len(indices)
    winners_in_circle = np.sum(is_winner[indices])
    win_pct = winners_in_circle / total_in_circle if total_in_circle > 0 else 0.0
    win_percentages.append(win_pct)

# Add results to dataframe
winnerDF['win_percentage'] = win_percentages

# Output results
for pct in win_percentages:
    print(pct)

Why This Is Better: #

• O(M log N) time complexity: Instead of checking every point for each winner, the KD-Tree only checks points in the spatial vicinity of the center. For 5000 winners and 320k rows, this cuts runtime from hours to minutes (or even seconds).
• Batch queries: query_ball_point processes all winner points in one call, reducing overhead.

Extra Tips to Speed Things Up More #

• Use itertuples() instead of iterrows(): For even faster row iteration over winnerDF, replace iterrows() with itertuples() (access fields like row.x instead of row['x']).
• Minimize IO operations: Collect all results in a list first, then print or save them all at once—printing inside loops adds unnecessary slowdown.
• Use dtype optimization: Specify data types when loading the CSV to reduce memory usage (e.g., dtype={'x': 'int32', 'y': 'int32', 'rank': 'int8'}).

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