如何无循环为二维数组按坐标与时间序列赋值？

需求说明 #

我们有包含坐标xs、ys，时间戳ts和极性ps的数据集，需要为图像每个(x,y)位置设置最新时间戳对应的极性值（xs与ys存在重复）。当前已通过循环实现功能，以下是向量化的替代方案：

向量化代码实现 #

import numpy as np

color_p = (0, 0, 255)
color_n = (255, 0, 0)
H, W = 128, 128
n_p = 100

xs, ys = np.random.randint(0, W, n_p), np.random.randint(0, H, n_p)
ts = np.random.rand(n_p)
ts.sort()  # 时间戳排序，确保后出现的元素时间更新
ps = np.random.randint(0, 2, n_p)

img = np.zeros((H, W, 3), dtype=np.uint8)

# 向量化核心步骤
# 1. 将(x,y)转换为唯一线性索引，简化去重逻辑
idx = ys * W + xs
# 2. 获取每个唯一坐标最后一次出现的索引（因ts已排序，最后出现即对应最新时间）
reverse_idx = idx[::-1]
_, reverse_unique_indices = np.unique(reverse_idx, return_index=True)
last_indices = n_p - 1 - reverse_unique_indices

# 3. 提取最新的坐标与极性数据
latest_xs = xs[last_indices]
latest_ys = ys[last_indices]
latest_ps = ps[last_indices]

# 4. 批量分配颜色并赋值到图像
img[latest_ys, latest_xs] = np.where(latest_ps[:, None], color_p, color_n)

代码解释 #

• 线性索引转换：通过ys * W + xs将二维坐标转为一维整数，避免直接处理二维坐标去重的复杂度。
• 获取最新索引：由于ts已排序，数组末尾的元素对应更新的时间戳。通过反转数组后调用np.unique获取首次出现的索引，再转换为原数组的末尾索引，即可得到每个(x,y)最后一次出现的位置。
• 批量赋值：用np.where根据极性批量生成匹配形状的颜色数组，直接通过索引完成图像赋值，完全替代循环操作，在数据量较大时能显著提升处理效率。

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