基于另一数组修改Python数组列值的高效NumPy实现求助

给定两个NumPy数组，其中arr_2的坐标（前两列）是arr_1的子集，具体如下：

import numpy as np

arr_1 = np.array([[2,5,25],
                  [2,6,28],
                  [7,6,75],
                  [8,4,39],
                  [3,1,86],
                  [1,4,45],
                  [6,5,89],
                  [3,4,99]])

arr_2 = np.array([[7,6,28],
                  [8,4,65],
                  [1,4,22],
                  [6,5,33]])

需求是：根据arr_2的(x,y)坐标（前两列），替换arr_1中对应位置的第三列（温度值），得到目标数组new_arr：

new_arr = np.array([[2,5,25],
                    [2,6,28],
                    [7,6,28],    # arr_1中的75替换为28
                    [8,4,65],    # arr_1中的39替换为65
                    [3,1,86],
                    [1,4,22],    # arr_1中的45替换为22 
                    [6,5,33],    # arr_1中的89替换为33
                    [3,4,99]])

原代码因嵌套循环导致大数据集下效率极低：

for i in range(len(self.dataset[key])):   
    ind = (np.array(self.pad["mesh_pos"][i])[:, None] == np.array(self.dataset["mesh_pos"][i])).all(-1).any(-1)
           
    id = np.array((ind==True).nonzero())[0]
    
    for k in range(len(id)):    
        ar[id[k],0] = self.dataset[key][i][0][k]  

    vr.append(ar)

self.dataset[key] = np.stack(vr, axis=0)

高效实现方案 #

核心思路是利用NumPy的向量化操作，彻底避免嵌套循环，大幅提升处理速度。以下提供两种适配不同场景的方案：

方案1：字典映射法（通用型）

通过将坐标转换为唯一键，构建坐标到温度值的映射，实现批量替换：

# 复制原数组避免修改原始数据
new_arr = arr_1.copy()

# 将(x,y)坐标转换为64位整数作为唯一标识（适用于整数坐标）
# 原理：把两列int32合并为一个int64，等价于 x * 2^32 + y，确保坐标唯一
keys_1 = arr_1[:, :2].view(np.int64).ravel()
keys_2 = arr_2[:, :2].view(np.int64).ravel()

# 构建坐标到温度的映射字典
temp_map = dict(zip(keys_2, arr_2[:, 2]))

# 筛选出arr_1中需要替换的位置
mask = np.isin(keys_1, keys_2)

# 批量替换对应位置的温度值
new_arr[mask, 2] = np.array([temp_map[key] for key in keys_1[mask]])

方案2：纯NumPy排序匹配法（高性能型）

如果坐标可以排序，用np.searchsorted实现纯向量化匹配，效率比字典法更高：

new_arr = arr_1.copy()

# 对两个数组的(x,y)坐标进行排序
sorted_idx_1 = np.lexsort(arr_1[:, :2].T)
sorted_arr_1 = arr_1[sorted_idx_1]
sorted_idx_2 = np.lexsort(arr_2[:, :2].T)
sorted_arr_2 = arr_2[sorted_idx_2]

# 查找x坐标匹配的位置，再验证y坐标是否一致
match_pos = np.searchsorted(sorted_arr_1[:, 0], sorted_arr_2[:, 0], side='left')
match_pos = match_pos[np.equal(sorted_arr_1[match_pos, 1], sorted_arr_2[:, 1])]

# 批量替换温度值
sorted_arr_1[match_pos, 2] = sorted_arr_2[:, 2]

# 恢复原数组的顺序
new_arr = sorted_arr_1[np.argsort(sorted_idx_1)]

方案优势 #

• 时间复杂度从原代码的O(n*m)降至O(n log n)或O(n)，大数据集下效率提升显著
• 两种方案适配不同场景：字典法对任意整数坐标友好，纯NumPy法在坐标有序时性能最优

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