Matlab中提取满足特定条件的三个向量对应元素位置的高效方法求助

适用场景 #

现有三个1×J大小的行向量A、B、C，需构造K×3大小的矩阵D，每行存储满足以下全部规则的三元组(a,b,c)：

• a为元素A(a)在向量A中的位置索引
• b为元素B(b)在向量B中的位置索引
• A(a)-B(b)的计算结果为向量C中的元素
• c为上述差值A(a)-B(b)在向量C中的位置索引
• A(a)与B(b)均不等于Inf、-Inf

实现代码 #

全程采用向量化运算，无显式循环，适合高频重复调用场景：

% 步骤1：过滤A、B中的无穷值，保留原始索引和对应数值
validA = ~isinf(A);
idxA = find(validA);
valA = A(validA);

validB = ~isinf(B);
idxB = find(validB);
valB = B(validB);

% 步骤2：向量化计算所有合法组合的差值
diff_mat = valA' - valB;

% 步骤3：匹配差值在C中的索引，筛选符合条件的组合
[match_flag, c_idx] = ismember(diff_mat, C);
[row, col] = find(match_flag);

% 步骤4：生成最终结果矩阵
D = [idxA(row), idxB(col), c_idx(match_flag)];

示例验证 #

输入题目给出的测试样例：

A = [-3 3 0 Inf -Inf];
B = [-2 2 0 Inf -Inf];
C = [Inf -Inf -1 1 0];

运行上述代码后得到的D和题目给出的参考结果完全一致：

D =
     1     1     3
     2     2     4
     3     3     5

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