Getting rid of nesting for loops which rely on reordering

I have two K x N matrices B and H.. I need to calculate a new matrix z according to:

     z(k,n) = a + B(k,n)*H(k,n) + ...
             sum{j = 1...(i-1)} 2B(j,n)*H(j,n)t +...
             sum{j = (i+1)...K)} B(j,n)*H(j,n)

where a and t are scalar constants and i depends on the "rank" in the column if H is sorted descending, i.e. for

     H = [1,2;5,6;3,4];

Column ranks are 2,3,1 for both columns. Sorry if this is confusing, I'm doing my best.

Anyways, I need to do this calculation for every element of z and I'm always only worried about the rank in each column and all terms in z are of elements of H and B in the same column.

My naive attempt to solve is using nested for loops:

    z = H.*B + a;
    for n=1:N       
        [Hsorted, idx] = sort(H(:,n),'descend');  % look at each column      
        for i=1:K                   %for each element of z(:,n)
            for j=1:K               %look at every other element in H(:,n)
                if j ~= i 
                    if  j < i                             
                        z(idx(i),n) = z(idx(i),n) + 2*B(idx(j),n)*Hsorted(j)*t;
                    else                               
                        z(idx(i),n) = z(idx(i),n) + B(idx(j),n)*Hsorted(j);
                    end
                end
            end  
        end 
    end

As you can probably imagine, this takes a lot of time for large K and N. I also am doing this as part of an iterative algorithm and a calculation like this gets done 2 or 3 times on each iteration. So, even for pretty small K and N, the time really adds up.

I've figured out that I can sort H, then reorder B according to the order of H with

     [Hsorted,orderH] = sort(H,'descend');
     [m,n] = size(B);
     Bsorted = B(sub2ind([m,n],orderH,repmat(1:n,m,1)));

and I'm guessing (hoping) that I can somehow replicate the ranked summation part of calculating the elements of z with matrix operations. This is where I'm stuck. Any help/advice or even a direction would be really appreciated

Answer 1

[Matt J]

      [K,N]=size(H);

       z = H.*B + a;
       [Hsorted, orderH] = sort(H,1,'descend');

       %sort B consistently with Hsorted
       idx=sub2ind([K,N],orderH,repmat(1:N,K,1));
       Bsorted=B(idx);

           e=ones(K);
           W=(triu(e,1)+2*t*tril(e,-1));

       z(idx)=z(idx)+W*(Bsorted.*Hsorted);

[Daniel Tweed]

Thanks for the quick response!

This is really helping me in the right direction, but I tested and get an error on the delta calculation if K and N are not square (as in the example H I posted with). Also, when I tested with H and K square, I get a different answer from your code than from my original. I really do need it to work with non-square matrices...

Tested with:

a = 0.5; t=0.25; H = B =  [1,2;5,6];

[Matt J]

Sorry, small mistake. I've edited my original answer now and see agreement with your code to within float precision. When I test with the following, I see a pretty decent speed-up

a = 0.5; t=0.25; 
H =   rand(1000,500);
B=H;

Elapsed time is 5.713132 seconds. %original version
Elapsed time is 0.115734 seconds. %new version

though I'm not sure what factor you were hoping for.

[Daniel Tweed]

That's a really good improvement. I wasn't expecting anything to great.

I actually see a much bigger improvement on my end (~300 times faster). I have several functions I need to calculate in the loop, and I think I can modify your answer to get those other functions out too. It's a HUGE help. Thanks!

[Daniel Tweed]

I posted a follow up question here and I don't know if it'd be too much to ask if you could take a look. If not, thanks for what you've already given me.