Assign values to a matrix based on the location of non-zero elements in another matrix

Question

I have a matrix, A containing ones and zeros. I want to create a new matrix B where the non-zero elements are in the same position as those in A, but where the values are given by the location of the first non-zero element in that row.

A few examples to show what I mean:

1 1 0 0 -> 1 1 0 0 (the first non-zero element is in column 1)
0 1 0 1 -> 0 2 0 2 (the first non-zero element is in column 2)
0 0 1 0 -> 0 0 3 0 (the first non-zero element is in column 3)
0 0 0 1 -> 0 0 0 4 (the first non-zero element is in column 4)

The code I have now is:

for i = 1:size(A,1)
for j = 1:size(A,2)
if A(i,j) == 1
B(i,:) = A(i,:) * find(A(i,:) == 1,1);
end
end
end

This gives the following results for a given A:

A
A =
 0 0 1 0
 0 1 1 1
 0 0 0 1
 1 1 1 1
 1 1 0 1
 0 1 1 1
 1 1 1 1
 1 1 1 1
 0 0 1 1
 0 0 1 1
B
B =
 0 0 3 0
 0 2 2 2
 0 0 0 4
 1 1 1 1
 1 1 0 1
 0 2 2 2
 1 1 1 1
 1 1 1 1
 0 0 3 3
 0 0 3 3

This code is slow if A is big. Is there a way to improve this code? I'm looking for improvements both on performance, and on the coding style / best practices etc.

Answer 1

First, let's go through the code:

You have two calls to size() in the beginning of your code, where the dimension is specified. It's better to make a single call to size, and save the two variables. This won't improve performance much, but will make it possible to reuse the two dimensions:

[rows, cols] = size(A);

---

You are using i and j as names for the iterators. In general, using those names are discouraged. It is not a problem when working with only real numbers, but once you start working with complex numbers this could be a problem. As an example, consider the following code, where we're adding complex numbers:

x = 1 + i; % Complex number, Real(x) = 1, Imag(x) = 1
for i = 1:10
 x = x + i
end

In this case, x will be 11 + i, instead of 10 + 11*i. Similarly, you might get code that doesn't behave as expected, if you forget to intialize i to a variable. For instance, if you think you have written i = true, i = false, i = 0 or something similar, neither of the two statements will give an error: i == true, i == false. This might be hard to debug since you don't get any error messages.

---

You have a growing loop inside your code. A row will be added to B for each iteration of the outer loop. This is very slow, and MATLAB actually warns you about this:

The Variable A appears to change size on every loop iteration

You should always preallocate memory. In this case, the size of the B matrix is known:

B = zeros(size(A));
% or
B = zeros(rows, cols);

---

You loop through all columns and use if A(i,j) == 1 to check if there are any non-zero elements. In the following code you use find to find the index of that element. A better way to do this is:

for ii = 1:rows
 idx = find(A(ii,:) == 1, 1);
 if ~isempty(idx)
 B(ii,:) = A(ii,:) * idx;
 end
end

This way, you can save one loop, and only loop through each element once.

---

You should use proper indentation. This can easily be achieved in Matlab. Press Ctrl+a followed by Ctrl+i, and Matlab will do the indentation for you.

---

The better way to do this:

A better way to do this, that's both simpler and faster is:

% Find the column index of the first element in each "slice".
[~, idx] = max(A,[],2); 
% Multiply the column index with each row of the initial matrix
bsxfun(@times, A, idx);

[M, idx] = max(A, [], dim) returns the value and index of the largest element along one dimension. It will only give one index per row or column (dependent on which dimension you choose). We can now multiply the vector of indices by the elements in A. For this, we use bsxfun.