超大型二维数组中，如何快速查找满足坐标条件的索引值？

核心优化思路 #

你的代码慢的根源有两个：每次调用都重复读取Excel（IO耗时极大）、用numpy数组的in操作做线性查找（时间复杂度O(n)，大型数组单次就需1秒）。针对这两个问题，以下是具体解决方法：

1. 缓存Excel数据，彻底消除重复IO #

用functools.lru_cache缓存读取和预处理后的结果，确保Excel只被读取一次，后续调用直接复用缓存数据：

import pandas as pd
import numpy as np
from functools import lru_cache
from scipy.interpolate import RegularGridInterpolator

@lru_cache(maxsize=None)
def get_preprocessed_grid():
    # 仅执行一次Excel读取
    df = pd.read_excel('data.xlsx')
    
    # 提取并预处理坐标与函数值
    x_unique = np.unique(df.iloc[:, 0].round(6))  # 处理浮点数精度，保留6位小数
    y_unique = np.unique(df.iloc[:, 1].round(6))
    value_grid = df.iloc[:, 2].values.reshape(len(x_unique), len(y_unique))  # 假设是规则网格
    
    # 构建O(1)查找的字典：(x,y) → 函数值
    grid_dict = {(x, y): value_grid[i, j] 
                 for i, x in enumerate(x_unique) 
                 for j, y in enumerate(y_unique)}
    
    # 预初始化插值器（重复调用无需重建）
    interpolator = RegularGridInterpolator((x_unique, y_unique), value_grid, method='linear')
    
    return grid_dict, interpolator

2. 用字典快速查找，替换numpy线性查找 #

修改你的核心函数，直接用缓存的字典做O(1)查找，插值逻辑也复用预初始化的插值器：

def assign_to_new_grid(p):
    grid_dict, interpolator = get_preprocessed_grid()
    
    # 对输入坐标做同样的精度处理
    x_key = round(p.x, 6)
    y_key = round(p.y, 6)
    
    if (x_key, y_key) in grid_dict:
        # 直接返回匹配的函数值
        return grid_dict[(x_key, y_key)]
    else:
        # 执行插值，直接调用预初始化的插值器
        return interpolator([[p.x, p.y]])[0]

关键细节说明 #

• 浮点数精度处理：用round(6)统一坐标精度，避免因浮点数存储误差导致的"明明坐标相同却匹配失败"问题；如果需要更灵活的公差控制，可替换为np.isclose判断：

  # 示例：用公差判断x是否匹配
  x_match = np.isclose(x_unique, p.x, atol=1e-6)
  y_match = np.isclose(y_unique, p.y, atol=1e-6)
  if np.any(x_match) and np.any(y_match):
      x_idx = np.where(x_match)[0][0]
      y_idx = np.where(y_match)[0][0]
      return value_grid[x_idx, y_idx]
  

• 规则网格假设：如果你的旧网格是不规则散点，可将RegularGridInterpolator替换为scipy.interpolate.griddata，同样预缓存散点坐标和值即可。

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