基于邻近位置填充Pandas DataFrame缺失值的优化需求

Problem Recap #

You have a Pandas DataFrame with a MultiIndex (location, date), where the x column has missing values. The goal is to fill these NaNs using values from neighboring locations at the same hour, following a specific priority order defined in the nearest DataFrame. The original iterative solution works but is too slow for large datasets—let's fix that.

Why the Original Code is Slow #

The original fillna_by_nearest function relies on nested Python loops:

• It iterates over every row with x.iteritems()
• For each missing value, it loops through each neighbor to find a valid entry
  This results in O(N*K) time complexity (N = number of rows, K = number of neighbors), which becomes painfully slow as your dataset scales. Vectorized Pandas operations (implemented under the hood in C) will drastically improve performance.

Optimized Solution #

We'll use a wide-table reshaping approach combined with vectorized filling to avoid Python-level loops entirely:

Step 1: Reshape to Wide Format

First, convert the long-form DataFrame to a wide format where each column represents a location, and each row represents an hour. This makes it trivial to access all neighbor values for a given time point:

import pandas as pd
import numpy as np

# Original data construction (from your example)
date = pd.date_range(start='2020-01-01', freq='H', periods=4)
locations = ["AA3", "AB1", "AD1", "AC0"]
x = [5.5, 10.2, np.nan, 2.3, 11.2, np.nan, 2.1, 4.0, 6.1, np.nan, 20.3, 11.3, 4.9, 15.2, 21.3, np.nan]
df = pd.DataFrame({'x': x})
df.index = pd.MultiIndex.from_product([locations, date], names=['location', 'date'])
df = df.sort_index()

# Neighbor priority definition
nearest = pd.DataFrame({
    "AA3": ["AA3", "AB1", "AD1", "AC0"],
    "AB1": ["AB1", "AA3", "AC0", "AD1"],
    "AD1": ["AD1", "AC0", "AB1", "AA3"],
    "AC0": ["AC0", "AD1", "AA3", "AB1"]
})

# Reshape to wide format (rows = dates, columns = locations)
wide_df = df.unstack(level='location')['x']

Step 2: Fill Missing Values with Vectorized Priority Logic

For each location, we'll use its neighbor priority list to create a sequence of columns, then use bfill(axis=1) to fill NaNs with the first valid value from the priority order:

filled_wide = pd.DataFrame()

for loc in wide_df.columns:
    # Get the priority list for the current location
    neighbor_order = nearest[loc].tolist()
    # Select columns in priority order, then fill NaNs by taking the first valid value from the sequence
    filled_col = wide_df[neighbor_order].bfill(axis=1).iloc[:, 0]
    filled_wide[loc] = filled_col

# Convert back to long-form DataFrame with the original MultiIndex
filled_df = filled_wide.stack().rename('x').swaplevel().sort_index()

Step 3: Verify the Result

Printing filled_df will give you exactly the output you expected:

x
location date          
AA3      2020-01-01 00:00:00   5.5
         2020-01-01 01:00:00  10.2
         2020-01-01 02:00:00   2.1
         2020-01-01 03:00:00   2.3
AB1      2020-01-01 00:00:00  11.2
         2020-01-01 01:00:00  10.2
         2020-01-01 02:00:00   2.1
         2020-01-01 03:00:00   4.0
AC0      2020-01-01 00:00:00   4.9
         2020-01-01 01:00:00  15.2
         2020-01-01 02:00:00  21.3
         2020-01-01 03:00:00  11.3
AD1      2020-01-01 00:00:00   6.1
         2020-01-01 01:00:00  15.2
         2020-01-01 02:00:00  20.3
         2020-01-01 03:00:00  11.3

Performance Benefits #

This approach leverages Pandas' optimized vectorized operations:

• No Python-level loops over rows or neighbors
• Reshaping and filling operations are handled in C, which is orders of magnitude faster for large datasets
• For a dataset with 10,000 rows, this method will run in seconds compared to minutes with the original loop-based code

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