PyTorch无循环构建矩阵B及获取矩阵C首零索引问题

一、构建矩阵B #

已知n×n的PyTorch张量A，元素为0到n-1的索引，需构建同形状张量B，使B[A[i,j]]=1（每行中A指定的列位置设为1），无需循环实现：

方法1：使用scatter_（推荐） #

scatter_是PyTorch中专门用于按索引赋值的方法，代码简洁高效：

import torch

# 示例：生成n×n的随机索引矩阵A
n = 5
A = torch.randint(0, n, (n, n))

# 初始化B为全0张量
B = torch.zeros((n, n), dtype=torch.int)
# 按A的索引将B对应位置设为1
B.scatter_(1, A, 1)

说明：scatter_的参数中，1表示按行维度进行索引，A是索引矩阵，1是要填充的值。该操作会将每行中A指定的列位置设为1，若同一行有重复索引，后出现的会覆盖先出现的。

方法2：使用index_put_ #

通过生成行、列索引对进行赋值，适合更灵活的索引场景：

B = torch.zeros((n, n), dtype=torch.int)
# 生成行索引：每行索引重复n次，对应A的每个元素的行位置
row_indices = torch.arange(n).repeat_interleave(n)
# 生成列索引：A展平后的所有元素
col_indices = A.flatten()
# 对指定位置赋值为1
B.index_put_((row_indices, col_indices), torch.ones_like(col_indices))

二、构建矩阵C #

C是1×n的张量，存储B中每行第一个0的索引：

# 生成B中0位置的掩码张量，0的位置为1，1的位置为0
zero_mask = (B == 0).int()
# 找到每行第一个1的索引（即原B中第一个0的位置），并转为1×n形状
C = zero_mask.argmax(dim=1).unsqueeze(0)

说明：argmax(dim=1)会返回每行第一个最大值的索引，这里zero_mask中0的位置标记为1，因此返回的索引就是该行第一个0的位置。若某行全为1，argmax会返回0，可根据业务需求添加边界处理逻辑。

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