Improving speed using logical indexing vectorisation of a 3d array

I have a 3-d array storing doubles named A. A has elements which can be split into two categories: values of A(:,:,1) being greater or equal to 0 or values of A(:,:,1) being less than zero. I would like to find the indexes of the later group of elements and set them equal to 0.

At the moment my code is a large "for" loop which I have learnt is inefficient in the MatLab context.

[m,n,p] = size(A);

for i = 1:m
    for j = 1:n
        if A(i,j,1) < 0
            A(i,j,1:3) = zeros();

During my attempts to speed up the code I have learnt about vectorization. However, I am struggling to see how to vectorize this problem. Any suggestion would be greatly appreciated.

Answer 1

[Jos (10584)]

So, let's see if I understood you correctly:

• A is an m-by-n-by-3 array to begin with
• If A(i,j,1) is less then 0, the three values in A(i,j,1:3) should be set to zero

A = 1+randi(5,[2 4]) ; 
A(:,:,2) = 10 * A(:,:,1) ;
A(:,:,3) = 10 * A(:,:,2)+1 ;
A(1,2,1) = -1  % some n-by-m-by3 array

% Unfortunately, the direct approach with logical indexing will fail
% tf = A(:,:,1) < 0
% A(tf,1:3) = 0 ; % alas, 3D arrays are tricky

[n,m,p] = size(A)  
A = reshape(A,[],p) % flatten it into a 2D array
tf = A(:,1) < 0
A(tf,1:3) = 0 
A = reshape(A,n,m,p)