arrays - 从矩阵的每一行中获取前 2 个非零元素

Question

我有一个A这样的矩阵：

A = [ 1 0 2 4; 2 3 1 0; 0 0 3 4 ]

A除了零之外，只有唯一的行元素，并且每行至少有 2 个非零元素。

我想从中创建一个新矩阵B，A其中的每一行都B包含 中相应行的前两个非零元素A。

B = [ 1 2 ; 2 3 ; 3 4 ]

循环很容易，但我需要矢量化解决方案。

标签： arraysmatrixvectorizationoctave

Answer 1

这是另一种矢量化方法。

A_bool = A > 0;   A_size = size(A);   A_rows = A_size(1);
A_boolsum = cumsum( A_bool, 2 ) .* A_bool;   % for each row, and at each column,
                                             % count how many nonzero instances
                                             % have occurred up to that column
                                             % (inclusive), and then 'zero' back
                                             % all original zero locations.

[~, ColumnsOfFirsts  ] = max( A_boolsum == 1, [], 2 );
[~, ColumnsOfSeconds ] = max( A_boolsum == 2, [], 2 );

LinearIndicesOfFirsts  = sub2ind( A_size, [1 : A_rows].', ColumnsOfFirsts  );
LinearIndicesOfSeconds = sub2ind( A_size, [1 : A_rows].', ColumnsOfSeconds );

Firsts  = A(LinearIndicesOfFirsts );
Seconds = A(LinearIndicesOfSeconds);

Result = horzcat( Firsts, Seconds )
% Result = 
%    1   2
%    2   3
%    3   4

---

_{PS。Matlab / Octave 通用子集兼容代码。}