如何优化Pandas DataFrame中重叠矩形的查找性能？

问题背景 #

我有一个包含数千组矩形的Pandas DataFrame，列包括group_id、x_min、x_max、y_min、y_max。目标是提取每组中至少与另一个矩形相交的所有矩形（相交判定：x和y方向均存在重叠）。现有基础代码可运行但速度过慢，希望做向量化优化。

基础实现代码 #

import pandas as pd

def find_overlapping_cuboids(df, check_z_plane=False):
    df_overlap = pd.DataFrame(columns=df.columns)
    group_ids = df["group_ID"].unique()

    for group_id in group_ids:
        df_filtered = df[df["group_ID"] == group_id]
        for i, row1 in df_filtered.iterrows():
            for j, row2 in df_filtered.iterrows():
                if i <= j:
                    continue  # Skip self-comparison
                if (
                    row1["x_min"] < row2["x_max"] and row1["x_max"] > row2["x_min"]
                ):  # Check for overlap in x direction
                    if (
                        row1["y_min"] < row2["y_max"] and row1["y_max"] > row2["y_min"]
                    ):  # Check for overlap in y direction
                        if not check_z_plane or (
                            row1["z_min"] < row2["z_max"]
                            and row1["z_max"] > row2["z_min"]
                        ):  # Check for overlap in z direction (optional)
                            df_overlap.loc[len(df_overlap)] = row1
                            df_overlap.loc[len(df_overlap)] = row2
    return df_overlap.drop_duplicates()

测试用DataFrame #

df = pd.DataFrame(
    {
        "group_ID": [1, 1, 1, 2, 2, 1],
        "x_min": [5, 0, 11, 4, 0, 16],
        "x_max": [15, 6, 16, 5, 4, 20],
        "y_min": [5, 0, 11, 1, 1.5, 16],
        "y_max": [15, 6, 16, 5, 4, 20],
        "z_min": [5, 0, 11, 4, 0, 16],
        "z_max": [15, 10, 16, 5, 5, 20],
    }
)

df_overlap = find_overlapping_cuboids(df, check_z_plane=False)

print(df_overlap)

预期输出 #

预期输出为包含以下行的DataFrame（对应原测试数据中group 1的前两个矩形、group 2的两个矩形）：

已尝试的向量化优化代码 #

import itertools
import numpy as np
import pandas as pd

def list_of_tuples_to_series(tuples, item_num):
    # this function takes a list of tuples and returns 1 series
    # the first digit in the tuple starts with item_num = 0
    # tuples = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
    # test = list_of_tuples_to_series(tuples,2)
    series = pd.Series([tuple[item_num] for tuple in tuples])
    return series


def find_overlapping_cuboids_fast(df, check_z_plane=True):
    df_overlap = pd.DataFrame(columns=["overlap_ID"] + list(df.columns))
    group_ids = df["group_ID"].unique()  # Get unique group_IDs
    overlap_id = 1  # Initialize overlap ID counter

    for group_id in group_ids:
        print(f"starting group {group_id}")
        df_filtered = df[
            df["group_ID"] == group_id
        ]  # Filter the dataframe for the current group_ID
        x_min = df_filtered["x_min"].values  # Extract the columns as NumPy arrays
        x_max = df_filtered["x_max"].values
        y_min = df_filtered["y_min"].values
        y_max = df_filtered["y_max"].values
        # Z = df_filtered['Z'].values
        # z_max = df_filtered['z_max'].values

        indices = np.arange(
            len(x_min)
        )  # Create all possible pairs of indices for comparison
        pairs = itertools.combinations(indices, 2)
        pairs = [
            pair for pair in pairs if pair[0] != pair[1]
        ]  # Exclude self-comparisons and repeated pairs
        i = list_of_tuples_to_series(pairs, 0)
        j = list_of_tuples_to_series(pairs, 1)

        overlap_x = (x_min[i] < x_max[j]) & (
            x_max[i] > x_min[j]
        )  # Check for overlap in x direction
        overlap_y = (y_min[i] < y_max[j]) & (
            y_max[i] > y_min[j]
        )  # Check for overlap in y direction
        # overlap_z = (~check_z_plane) | ((Z[i] < z_max[j]) & (z_max[i] > Z[j]))# Check for overlap in z direction (optional)
        overlap_conditions = (
            overlap_x & overlap_y
        )  # & overlap_z # Combine all overlap conditions
        overlapping_indices = np.where(overlap_conditions)[
            0
        ]  # Find the indices of overlapping cuboids

        for idx in overlapping_indices:
            row1 = df_filtered.iloc[i[idx]]
            row2 = df_filtered.iloc[j[idx]]
            row1_with_id = pd.Series(
                [overlap_id] + list(row1), index=df_overlap.columns
            )
            row2_with_id = pd.Series(
                [overlap_id] + list(row2), index=df_overlap.columns
            )
            df_overlap.loc[len(df_overlap)] = row1_with_id
            df_overlap.loc[len(df_overlap)] = row2_with_id
            overlap_id += 1  # Increment overlap ID
    return df_overlap.drop_duplicates()


# Find overlapping cuboids
df_overlap_fast2 = find_overlapping_cuboids_fast(df, check_z_plane=False)

进一步提速建议 #

1. 完全向量化配对比较，避免迭代组合 #

使用NumPy的广播机制生成所有i<j的配对矩阵，替代itertools.combinations的迭代，大幅减少循环开销：

def find_overlapping_cuboids_vectorized(df, check_z_plane=False):
    result_dfs = []
    overlap_id = 1
    
    for group_id in df["group_ID"].unique():
        df_g = df[df["group_ID"] == group_id].reset_index(drop=True)
        n = len(df_g)
        if n < 2:
            continue
        
        # 生成i < j的布尔矩阵
        i, j = np.triu_indices(n, k=1)
        
        # 向量化计算重叠条件
        x_overlap = (df_g["x_min"].values[i] < df_g["x_max"].values[j]) & (df_g["x_max"].values[i] > df_g["x_min"].values[j])
        y_overlap = (df_g["y_min"].values[i] < df_g["y_max"].values[j]) & (df_g["y_max"].values[i] > df_g["y_min"].values[j])
        
        if check_z_plane:
            z_overlap = (df_g["z_min"].values[i] < df_g["z_max"].values[j]) & (df_g["z_max"].values[i] > df_g["z_min"].values[j])
            overlap_mask = x_overlap & y_overlap & z_overlap
        else:
            overlap_mask = x_overlap & y_overlap
        
        # 获取所有重叠配对的索引
        valid_i = i[overlap_mask]
        valid_j = j[overlap_mask]
        
        # 批量构造结果，避免逐行添加
        if len(valid_i) > 0:
            # 提取配对的行
            rows_i = df_g.iloc[valid_i].assign(overlap_ID=overlap_id)
            rows_j = df_g.iloc[valid_j].assign(overlap_ID=overlap_id)
            
            # 合并并更新overlap_id
            group_result = pd.concat([rows_i, rows_j], ignore_index=True)
            result_dfs.append(group_result)
            overlap_id += len(valid_i)
    
    # 合并所有组结果并去重
    final_df = pd.concat(result_dfs, ignore_index=True).drop_duplicates()
    # 调整列顺序，把overlap_ID放前面
    final_df = final_df[["overlap_ID"] + [col for col in df.columns if col != "overlap_ID"]]
    return final_df

2. 避免逐行添加DataFrame #

原代码中df_overlap.loc[len(df_overlap)] = ...是非常慢的操作，改用pd.concat批量合并DataFrame片段，或者预先构造数组再转换为DataFrame。

3. 利用空间索引优化（针对超大规模数据） #

如果每组的矩形数量非常多（比如上万级），可以使用rtree库构建空间索引，快速查询相邻矩形，减少不必要的配对比较：

pip install rtree

from rtree import index

def find_overlapping_with_rtree(df, check_z_plane=False):
    result_dfs = []
    overlap_id = 1
    
    for group_id in df["group_ID"].unique():
        df_g = df[df["group_ID"] == group_id].reset_index(drop=True)
        n = len(df_g)
        if n < 2:
            continue
        
        # 构建RTree索引
        idx = index.Index()
        for pos, row in df_g.iterrows():
            # 索引存储矩形的边界 (x_min, y_min, x_max, y_max)
            idx.insert(pos, (row["x_min"], row["y_min"], row["x_max"], row["y_max"]))
        
        # 存储已处理的配对，避免重复
        processed = set()
        group_pairs = []
        
        for pos, row in df_g.iterrows():
            # 查询与当前矩形重叠的所有矩形索引
            hits = list(idx.intersection((row["x_min"], row["y_min"], row["x_max"], row["y_max"])))
            for hit in hits:
                if hit <= pos or (pos, hit) in processed:
                    continue
                processed.add((pos, hit))
                # 验证重叠条件（RTree的intersection可能返回边界接触的，需二次确认）
                other_row = df_g.iloc[hit]
                x_ok = row["x_min"] < other_row["x_max"] and row["x_max"] > other_row["x_min"]
                y_ok = row["y_min"] < other_row["y_max"] and row["y_max"] > other_row["y_min"]
                z_ok = not check_z_plane or (row["z_min"] < other_row["z_max"] and row["z_max"] > other_row["z_min"])
                if x_ok and y_ok and z_ok:
                    group_pairs.append((pos, hit))
        
        # 批量生成结果
        if group_pairs:
            i_list, j_list = zip(*group_pairs)
            rows_i = df_g.iloc[list(i_list)].assign(overlap_ID=overlap_id)
            rows_j = df_g.iloc[list(j_list)].assign(overlap_ID=overlap_id)
            group_result = pd.concat([rows_i, rows_j], ignore_index=True)
            result_dfs.append(group_result)
            overlap_id += len(group_pairs)
    
    final_df = pd.concat(result_dfs, ignore_index=True).drop_duplicates()
    final_df = final_df[["overlap_ID"] + [col for col in df.columns if col != "overlap_ID"]]
    return final_df

4. 其他小优化 #

• 提前将DataFrame的列转换为NumPy数组，减少重复访问开销；
• 去掉不必要的print语句，减少IO开销；
• 使用reset_index(drop=True)确保分组后的索引连续，避免iloc的潜在问题。

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