MATLAB中矩阵重复元素列索引的向量化求解方法

现有N×3的正整数矩阵A，需生成尺寸为max(max(A))×2的矩阵Q，规则如下：

• Q的第i行对应A中值为i的元素的列索引
• 若A中值i出现两次，Q(i,:)为这两个位置的列索引
• 若A中值i仅出现一次，Q(i,:)的两列值相同

原循环实现效率低下，当A行数多、max(max(A))值大时尤为明显，以下是向量化解决方案。

示例输入 #

A = [  1    10     9
       1     2     4
       3     2    11
       3    12     5
       6    13     4
       6     7    15
       8     7     5
       8    14    16];

预期输出 #

Q =
     1     1
     2     2
     1     1
     3     3
     3     3
     1     1
     2     2
     1     1
     3     3
     2     2
     3     3
     2     2
     2     2
     2     2
     3     3
     3     3

原低效循环代码 #

max_Q = max(max(A));
Q     = zeros(max_Q,2);
for i = 1 : max_Q
    [~,cols] = find(i == A); 
    Q(i,:) = cols;
end

向量化实现代码 #

% 获取A中的最大值，确定Q的行数
max_val = max(A(:));
% 将A展平为一维数组，对应每个元素的值
vals = A(:);
% 生成每个元素对应的列索引，展平为一维数组
cols = repmat(1:3, size(A,1), 1);
cols = cols(:);
% 使用accumarray按值分组，提取每组的第一个和最后一个列索引
% 若值仅出现一次，第一个和最后一个索引相同，自动满足两列重复的要求
Q = accumarray(vals, cols, [max_val, 2], @(x) [x(1), x(end)]);

代码解释 #

1. 展平矩阵与列索引生成：将A展平为一维数组vals，同时生成每个元素对应的列索引数组cols（因为A是N×3结构，列索引固定为1、2、3循环重复）。
2. accumarray分组处理：accumarray会自动按vals中的值对cols进行分组，对每个分组应用匿名函数@(x) [x(1), x(end)]：
   • 当值i在A中出现两次时，分组x包含两个列索引，取第一个和最后一个填充Q的第i行
   • 当值i在A中仅出现一次时，分组x只有一个索引，x(1)和x(end)为同一值，自动实现两列重复的要求
3. 自动填充Q：accumarray的第三个参数[max_val, 2]直接指定输出矩阵的尺寸，无需额外初始化，效率远高于循环遍历每个i值。

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