CODEWITHSUNDEEP
matlabvectorization
KevinKim
KevinKim

Reputation: 1434

return the index of the last K non-zero element of each row of a matrix

Is there a vectorization way of returning the index of the last K nonzero elements of each row of a matrix?

For example, my matrix only contains 0 and 1 and the last column of each row is always 1. Then I want to find the index of the last K, where K>1, nonzero elements of each row. If a row only has M (less than K) nonzero elements, then the index for that row is just the index of the last M nonzero element. e.g.

A = [0 1 0 1;
     1 1 0 1;
     1 1 1 1;
     0 0 0 1]

And my K = 2, then I expected to return a matrix such that

B = [0 1 0 1;
     0 1 0 1;
     0 0 1 1;
     0 0 0 1]

Namely B is originally a zero matrix with same shape as A, then it copies each row of A where the corresponding column starts from the index of the last K non-zero element of the row of A (and if in one row of A there is only M < K non-zero element, then it starts from the index of the last M non-zero element of that row of A)

Upvotes: 0

Views: 240

Answers (2)

rahnema1
rahnema1

Reputation: 15867

knowing that find function can get indices of last k elements, we can use bsxfun to apply find to rows of a matrix to find which element in each row satisfy the condition. find again used to extract rows and columns of nonzero elements of the resultant matrix, so reducing size of data and complexity of operations. then save the result to a sparse matrix then convert to full matrix:

A = [0 1 0 1;
     1 1 0 1;
     1 1 1 1;
     0 0 0 1]
k = 2;
[row , col]= size(A);
last_nz = bsxfun(@(a,b)find(a,b,'last'),A',(repmat(k, 1, row))); %get indices of last two nonzero elements for each row
[~,rr,cc]=find(last_nz); %get columns and rows of correspondong element for whole matrix
B = full(sparse(rr,cc,1));

Upvotes: 0

beaker
beaker

Reputation: 16810

Knowing that elements are only 0 or 1, you can make a mask using cumsum on the flipped matrix A and throw away values with a cumulative sum greater than k: 

A = [0 1 0 1;1 1 0 1;1 1 1 1;0 0 0 1]
k = 2;

C = fliplr(cumsum(fliplr(A), 2));   % take the cumulative sum backwards across rows
M = (C <= k);                       % cumsum <= k includes 0 elements too, so...
B = A .* M                          % multiply original matrix by mask

As mentioned in the comments (Thanks @KQS!), if you're using a recent version of MATLAB, there's a direction optional parameter to cumsum, so the line to generate C can be shortened to:

C = cumsum(A, 2, 'reverse');

Results: A =

   0   1   0   1
   1   1   0   1
   1   1   1   1
   0   0   0   1

B =

   0   1   0   1
   0   1   0   1
   0   0   1   1
   0   0   0   1

Upvotes: 2

Related Questions