如何基于最小欧氏距离选取行补全NYC MVA数据集缺失ZIP码

Hey there! Let's fix that slow imputation code for your NYC MVA dataset—your idea to use geospatial nearest neighbors is solid, but that apply approach was never going to handle 1.7 million rows efficiently. Here's a much faster way to get the job done:

Step 1: Clean Up the ZIP CODE Column First #

First, let's fix the messy ZIP CODE values you saw (mix of strings, floats, and empty spaces) to ensure our reference data is consistent:

import pandas as pd
import numpy as np

# Clean ZIP CODE formatting
df['ZIP CODE'] = df['ZIP CODE'].astype(str).str.strip()
# Convert entries like "10000.0" to "10000"
df['ZIP CODE'] = df['ZIP CODE'].str.replace('.0', '', regex=False)
# Replace empty strings with NaN for consistent handling
df['ZIP CODE'] = df['ZIP CODE'].replace('', np.nan)

Step 2: Split Data into Reference and Missing Groups #

Separate the rows that have valid ZIP codes (our reference points) from those that need imputation:

# Rows with valid ZIP codes
has_zip = df.dropna(subset=['ZIP CODE']).copy()
# Rows missing ZIP codes
missing_zip = df[df['ZIP CODE'].isna()].copy()

Step 3: Use KDTree for Fast Nearest Neighbor Search #

Instead of scanning the entire dataset for every missing row, we'll use sklearn's KDTree—a data structure optimized for fast spatial queries. This cuts the time complexity from O(n²) (your original approach) to O(n log n), which works seamlessly for large datasets:

from sklearn.neighbors import KDTree

# Extract latitude/longitude from reference data
reference_coords = has_zip[['LATITUDE', 'LONGITUDE']].values
# Build the KDTree
kdtree = KDTree(reference_coords, leaf_size=30)

# Extract coordinates from rows missing ZIP codes
missing_coords = missing_zip[['LATITUDE', 'LONGITUDE']].values
# Find the single nearest neighbor (k=1) for each missing point
distances, nearest_indices = kdtree.query(missing_coords, k=1)

# Map the nearest ZIP code to the missing rows
missing_zip['ZIP CODE'] = has_zip.iloc[nearest_indices.flatten()]['ZIP CODE'].values

# Optional: Add a column to track the distance to the nearest reference point
missing_zip['NN_DISTANCE'] = distances.flatten()

Step 4: Combine Data Back Together #

Merge the imputed rows with the original reference rows, preserving the original index order:

# Combine and re-sort to match original dataframe order
df_imputed = pd.concat([has_zip, missing_zip]).sort_index()

Why This Works So Much Better #

Your original apply method was checking every row in the dataset for each missing ZIP entry—with 500k+ missing rows, that's billions of operations. The KDTree pre-processes the reference coordinates into a spatial tree, so each nearest neighbor lookup takes logarithmic time instead of linear time.

Bonus Optimization Tips #

• Aggregate duplicate coordinates: If multiple rows share the same latitude/longitude, you can group them and keep one reference entry to reduce the size of the KDTree.
• Try BallTree instead: If you want to use spherical distance (more accurate for global data), BallTree is a good alternative—though for NYC's small geographic range, Euclidean distance works perfectly.
• Validate imputations: Spot-check some imputed rows by mapping the coordinates to ensure the ZIP code makes geographic sense.

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