Python中使用自定义坐标索引查找numpy二维数组取值的最优方法

核心需求是将自定义的非整数坐标映射到numpy原生的整数索引，无需依赖第三方库，直接用numpy内置方法即可实现高效优雅的效果。

实现逻辑 #

• 你定义的x、y坐标均为严格单调递增序列，直接用np.searchsorted即可快速返回目标坐标在自定义序列中对应的索引位置
• 增加浮点值接近度校验，避免浮点精度问题导致的取值错误
• 可封装为自定义类，实现和原生numpy数组一致的[]取值体验

代码示例 #

基础函数实现 #

import numpy as np

# 初始化示例数据
arr = np.array([[1,2,3,4], [5,6,7,8]])
custom_x = np.array([1, 3, 5, 7])
custom_y = np.array([1, 1.5])

def get_custom_val(target_x, target_y):
    # 映射x坐标到原生索引
    x_idx = np.searchsorted(custom_x, target_x)
    # 映射y坐标到原生索引
    y_idx = np.searchsorted(custom_y, target_y)
    # 坐标合法性校验
    if not np.isclose(custom_x[x_idx], target_x) or not np.isclose(custom_y[y_idx], target_y):
        raise ValueError(f"坐标({target_x}, {target_y})不在预设的自定义坐标范围内")
    return arr[y_idx, x_idx]

# 测试用例
print(get_custom_val(7, 1.5))
# 输出：8

更优雅的类封装实现 #

class CustomCoordArray:
    def __init__(self, raw_arr, x_coords, y_coords):
        self.raw_arr = np.asarray(raw_arr)
        self.x_coords = np.asarray(x_coords)
        self.y_coords = np.asarray(y_coords)
        # 校验坐标长度和数组维度匹配
        assert self.x_coords.shape[0] == self.raw_arr.shape[1]
        assert self.y_coords.shape[0] == self.raw_arr.shape[0]
    
    def __getitem__(self, coords):
        target_x, target_y = coords
        x_idx = np.searchsorted(self.x_coords, target_x)
        y_idx = np.searchsorted(self.y_coords, target_y)
        if not np.isclose(self.x_coords[x_idx], target_x) or not np.isclose(self.y_coords[y_idx], target_y):
            raise ValueError(f"坐标({target_x}, {target_y})不在预设范围内")
        return self.raw_arr[y_idx, x_idx]

# 使用示例
custom_arr = CustomCoordArray(arr, custom_x, custom_y)
print(custom_arr[7, 1.5])
# 输出：8

其他场景适配 #

• 如果自定义坐标不是单调序列，可提前构建坐标到索引的字典映射，适合坐标数量不大的场景，查询效率为O(1)：

  x_map = {val:i for i, val in enumerate(custom_x)}
  y_map = {val:i for i, val in enumerate(custom_y)}
  

  取值时直接调用arr[y_map[target_y], x_map[target_x]]即可。

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