vectorization of for loop

function o_lmat = reassign_label(o_lmat,h,w,p,n_prev)

for a = 1:h % for loops to reassign labels

for b = 1:w

for c = 1:p

if (o_lmat(a,b,c) ~= 0)

o_lmat(a,b,c) = o_lmat(a,b,c)+n_prev;

end

end

end

end

end

I am trying to vectorize this code as:

function o_lmat = reassign_label(o_lmat,h,w,p,n_prev)

% o_lmat = zeros(h,w,p);

if (o_lmat(:,:,:) ~= 0)

o_lmat(:,:,:) = o_lmat(:,:,:)+n_prev;

end

end

It is not working correctly. Pllease find the mistake I am making.

Answer 1

[Chunru]

function o_lmat = reassign_label(o_lmat,h,w,p,n_prev)
% for a = 1:h   % for loops to reassign labels
%     for b = 1:w
%         for c = 1:p
%             if (o_lmat(a,b,c) ~= 0)
%                 o_lmat(a,b,c) = o_lmat(a,b,c)+n_prev;
%             end
%         end
%     end
% end
idx = o_lmat(:) ~=0;
o_lmat(idx) = o_lmat(idx) + n_prev;
end

[Rizwan Khan]

thanks a lot