PyTorch张量值替换：无循环实现指定元素替换

要实现不使用循环，将张量A中所有与before匹配的元素替换为after中对应位置的元素，可以利用PyTorch的广播机制和向量化操作完成，以下是两种可行方案：

方案一：通用匹配替换（无需排序） #

这种方法不要求before张量有序，适用于任意场景：

import torch

def replace(A, before, after):
    # 克隆原张量，避免修改输入的A
    A_copy = A.clone()
    # 通过广播生成匹配矩阵：每个A元素与before所有元素逐一比较
    match_mask = A_copy[..., None] == before
    # 获取所有匹配位置对应的before元素索引
    match_indices = match_mask.nonzero(as_tuple=True)[1]
    # 筛选出A中需要替换的位置，并用after对应元素替换
    A_copy[match_mask.any(dim=-1)] = after[match_indices]
    return A_copy

# 测试示例
before = torch.Tensor([2,4,5])
after  = torch.Tensor([20,40,50])
A = torch.Tensor([1,2,3,4,5,6])

result = replace(A, before, after)
print(result)  # 输出: tensor([ 1., 20.,  3., 40., 50.,  6.])

核心逻辑 #

1. 克隆A避免直接修改原始输入；
2. 利用广播特性生成布尔矩阵，标记A中每个元素是否匹配before里的任意元素；
3. 通过nonzero提取匹配位置对应的before元素索引；
4. 用布尔掩码筛选出需要替换的位置，完成元素替换。

方案二：基于排序的高效替换 #

如果before可以提前排序，这种方法效率更高，适合处理大规模张量：

import torch

def replace_sorted(A, before, after):
    # 对before和对应的after按before元素排序
    sorted_idx = torch.argsort(before)
    sorted_before = before[sorted_idx]
    sorted_after = after[sorted_idx]
    
    # 查找A中元素在排序后before中的插入位置
    search_idx = torch.searchsorted(sorted_before, A)
    # 限制索引范围，防止越界
    search_idx = torch.clamp(search_idx, 0, len(sorted_before)-1)
    # 生成匹配掩码
    match_mask = sorted_before[search_idx] == A
    # 用torch.where完成替换
    return torch.where(match_mask, sorted_after[search_idx], A)

# 测试示例
result = replace_sorted(A, before, after)
print(result)  # 输出: tensor([ 1., 20.,  3., 40., 50.,  6.])

核心逻辑 #

1. 对before和after按before元素排序，保证before有序；
2. 用torch.searchsorted快速定位A元素在排序后before中的位置；
3. 通过比较生成匹配掩码，最后用torch.where完成替换。

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