PyTorch中查找张量B元素在张量A中所有索引位置的内存优化方案

import torch

def vectorized_find_indices_broadcast(A, B):
    # 扩展A的维度以与B进行广播比较
    # A_expanded = A[:, None, None] # 原始问题中的三重扩展可能并非必需，但原理相同
    # mask = (B == A_expanded)
    # ... 后续操作
    pass

import torch

def find_indices_hybrid(a: torch.Tensor, b: torch.Tensor):
    """
    使用部分广播和Python循环查找张量B中元素在张量A中的所有索引。

    Args:
        a (torch.Tensor): 目标张量A。
        b (torch.Tensor): 待查找值的张量B。

    Returns:
        list[list[int]]: 一个列表的列表，其中每个子列表包含B中对应值在A中的索引。
    """
    if a.ndim != 1 or b.ndim != 1:
        raise ValueError("输入张量a和b必须是一维张量。")

    # 步骤1: 使用unsqueeze和nonzero获取所有匹配的(a_idx, b_idx)对
    # (a.unsqueeze(1) == b) 会创建一个 N x M 的布尔张量，
    # 其中 (i, j) 为 True 表示 a[i] == b[j]。
    # nonzero() 将其转换为一个 K x 2 的张量，K是匹配的总数，
    # 每行 [a_idx, b_idx] 表示 a[a_idx] == b[b_idx]。
    overlap_idxs = (a.unsqueeze(1) == b).nonzero()

    # 步骤2: 初始化结果列表，为B中的每个元素预留一个子列表
    output = [[] for _ in b]

    # 步骤3: 遍历匹配对，将A的索引添加到B对应值的列表中
    for (a_idx, b_idx) in overlap_idxs:
        output[b_idx.item()].append(a_idx.item())

    return output

# 示例
A_tensor = torch.tensor([1, 2, 3, 3, 2, 1, 4, 5, 9])
B_tensor = torch.tensor([1, 2, 3, 9])
result_hybrid = find_indices_hybrid(A_tensor, B_tensor)
print(f"混合方法结果: {result_hybrid}")
# 预期输出: [[0, 5], [1, 4], [2, 3], [8]]

import torch

def find_indices_pure_loop(a: torch.Tensor, b: torch.Tensor):
    """
    使用纯Python循环查找张量B中元素在张量A中的所有索引。

    Args:
        a (torch.Tensor): 目标张量A。
        b (torch.Tensor): 待查找值的张量B。

    Returns:
        list[list[int]]: 一个列表的列表，其中每个子列表包含B中对应值在A中的索引。
    """
    if a.ndim != 1 or b.ndim != 1:
        raise ValueError("输入张量a和b必须是一维张量。")

    output = []
    for _b_val in b:
        # 对于B中的每个值，在A中查找其所有索引
        # (a == _b_val) 会生成一个 N 长度的布尔张量
        # .nonzero() 找到所有为True的索引
        # .squeeze() 移除不必要的维度（例如，如果只有一个索引，结果是(1,)而不是(1,1)）
        # .tolist() 转换为Python列表
        idxs = (a == _b_val).nonzero().squeeze().tolist()

        # 确保结果是列表形式，即使只有一个或没有匹配
        if not isinstance(idxs, list):
            idxs = [idxs] # 如果只有一个匹配，squeeze().tolist()可能返回一个int

        output.append(idxs)

    return output

# 示例
A_tensor = torch.tensor([1, 2, 3, 3, 2, 1, 4, 5, 9])
B_tensor = torch.tensor([1, 2, 3, 9])
result_pure_loop = find_indices_pure_loop(A_tensor, B_tensor)
print(f"纯循环方法结果: {result_pure_loop}")
# 预期输出: [[0, 5], [1, 4], [2, 3], [8]]