如何用Python在PyTorch张量中随机替换值且保留每个值至少一次？

需求说明 #

需要将张量中出现次数超过2次的值，部分替换为仅出现1次的值，替换后必须保证所有值的出现次数≥2次。

示例输入输出 #

输入张量：

edge_index = torch.as_tensor([[0, 0, 1, 2, 3, 4, 6, 7, 7, 8],
                              [1, 2, 2, 4, 4, 5, 0, 1, 3, 7]])

参考输出（结果不唯一，满足要求即可）：

result = torch.as_tensor([[5, 0, 1, 2, 3, 8, 6, 7, 7, 8],
                          [1, 2, 2, 4, 4, 5, 0, 1, 3, 6]])

原有代码问题 #

原代码通过随机选择高频值的位置替换为低频值，但无法保证替换后：

• 原高频值的剩余出现次数≥2次
• 不会出现原高频值被完全替换消失的情况

比如原代码运行后出现1消失、7仅出现1次的问题。

解决方案代码 #

import torch

edge_index = torch.as_tensor([[0, 0, 1, 2, 3, 4, 6, 7, 7, 8],
                              [1, 2, 2, 4, 4, 5, 0, 1, 3, 7]])
result = edge_index.clone()

# 统计每个节点的出现次数
n_id, n_counts = torch.unique(result, return_counts=True)
count_map = {id.item(): cnt.item() for id, cnt in zip(n_id, n_counts)}

# 分离低频节点（仅出现1次）和高频节点（出现>2次）
low_freq_nodes = n_id[n_counts == 1].tolist()
high_freq_nodes = n_id[n_counts > 2].tolist()

# 逐个处理低频节点，替换高频节点的位置
for low_node in low_freq_nodes:
    # 随机选择一个仍可被替换的高频节点（替换后次数≥2）
    high_node = torch.tensor(high_freq_nodes).random_().item()
    # 找到该高频节点在张量中的所有位置
    mask = (result == high_node)
    indices = torch.where(mask)
    # 随机选一个位置进行替换
    rand_pos_idx = torch.randint(0, len(indices[0]), (1,)).item()
    replace_pos = (indices[0][rand_pos_idx], indices[1][rand_pos_idx])
    result[replace_pos] = low_node
    
    # 更新计数并维护高频节点列表
    count_map[high_node] -= 1
    count_map[low_node] += 1
    if count_map[high_node] == 2:
        high_freq_nodes.remove(high_node)

# 验证所有节点出现次数≥2
final_counts = torch.unique(result, return_counts=True)[1]
assert (final_counts >= 2).all(), "替换后存在出现次数不足2次的节点"

print("最终替换结果：")
print(result)

方案说明 #

1. 精准控制替换对象：仅从出现次数>2的高频节点中选择替换位置，确保替换后该高频节点次数仍≥2（原次数>2，减1后≥2）
2. 低频节点达标：每个低频节点仅需替换一次，即可将出现次数从1提升至2
3. 动态维护高频列表：当高频节点次数降至2时，从列表中移除，避免后续被选中导致次数不足

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