Alex
Alex

Reputation: 15708

Efficient way to collapse each row of a matrix based on a column index for each row

I would like to reduce the number of columns in a matrix row-wise by performing an action depending on a column index specified for each row. Given the following example data:

M = magic(4);
col_ind = [1; 3; 2; 4];

I would like to make a three column matrix, such that the first column contains the row sum of each row up to the specified column index, the second column contains the value of M in that row, specified by the column index, and the last column contains the row sum of the rest of the entries in that row, that is, produce the matrix M_out:

M_out = nan(4,3);
for i = 1:4
    M_out(i, :) = [sum(M(i, 1:col_ind(i)-1), 2), ...
        M(i, col_ind(i)), sum(M(i, col_ind(i)+1:end), 2)];

end

such that:

>> M_out

M_out =

     0    16    18
    16    10     8
     9     7    18
    33     1     0

Upvotes: 2

Views: 146

Answers (1)

Erfan
Erfan

Reputation: 1927

According to Stewie's comment, I will do it here using tril. The second column is straightforward using sub2ind. Then the third column becomes trivial:

T = tril(ones(size(M, 2)), -1);
M_out = zeros(size(M, 1), 3);
M_out(:, 1) = sum(T(col_ind,:) .* M, 2);
M_out(:, 2) = M(sub2ind(size(M), (1:size(M, 1)).', col_ind));
M_out(:, 3) = sum(M, 2) - sum(M_out, 2);

Or, a bit more efficient using cumsum:

cs = cumsum(M, 2);
M_out = zeros(size(M, 1), 3);
M_out(:, 2) = M(sub2ind(size(M), (1:size(M, 1)).', col_ind));
M_out(:, 1) = cs(sub2ind(size(cs), (1:size(cs, 1)).', col_ind)) - M_out(:, 2);
M_out(:, 3) = cs(:, end) - sum(M_out, 2);

Upvotes: 3

Related Questions