Accumarray函数在c++中的实现

arma::mat accumarray (arma::mat& subs, arma::vec& val, arma::rowvec& sz)
{
    arma::u32 ar = sz.col(0)(0);
    arma::u32 ac = sz.col(1)(0);

    arma::mat A; A.set_size(ar, ac);

    for (arma::u32 r = 0; r < ar; ++r)
    {
        for (arma::u32 c = 0; c < ac; ++c)
        {
            arma::uvec idx = arma::find(subs.col(0) == r &&
                                        subs.col(1) == c);
            if (!idx.is_empty())
                A(r, c) = arma::sum(val.elem(idx));
            else
                A(r, c) = 0;
        }
    }

    return A;
}

The sz input is a two columns vector that contain : num rows / num cols for the output matrix A. The subs matrix is a 2 columns with same num rows of val. Num rows of val is basically sz.rows by sz.cols.

The sz (size) input is not really mandatory and can be deduced easily by searching the max in subs columns.

arma::u32 sz_rows = arma::max(subs.col(0)) + 1;
arma::u32 sz_cols = arma::max(subs.col(1)) + 1;

or

arma::u32 sz_rows = arma::max(subs.col(0)) + 1;
arma::u32 sz_cols = val.n_elem / sz_rows;

the output matrix is now :

arma::mat A (sz_rows, sz_cols);

the accumarray function become :

arma::mat accumarray (arma::mat& subs, arma::vec& val)
{
    arma::u32 sz_rows = arma::max(subs.col(0)) + 1;
    arma::u32 sz_cols = arma::max(subs.col(1)) + 1;

    arma::mat A (sz_rows, sz_cols);

    for (arma::u32 r = 0; r < sz_rows; ++r)
    {
        for (arma::u32 c = 0; c < sz_cols; ++c)
        {
            arma::uvec idx = arma::find(subs.col(0) == r &&
                                        subs.col(1) == c);
            if (!idx.is_empty())
                A(r, c) = arma::sum(val.elem(idx));
            else
                A(r, c) = 0;
        }
    }

    return A;
}

For example :

arma::vec val = arma::regspace(101, 106);

arma::mat subs;
subs << 0 << 0 << arma::endr
     << 1 << 1 << arma::endr
     << 2 << 1 << arma::endr
     << 0 << 0 << arma::endr
     << 1 << 1 << arma::endr
     << 3 << 0 << arma::endr;

arma::mat A = accumarray (subs, val);

A.raw_print("A =");

Produce this result :

A = 
    205    0
      0  207
      0  103
    106    0

This example is found here : http://fr.mathworks.com/help/matlab/ref/accumarray.html?requestedDomain=www.mathworks.com except for the indices of subs, armadillo is 0-based indice where matlab is 1-based.

Unfortunaly, the previous code is not suitable for big matrix. Two for-loop with a find in vector in between is really bad thing. The code is good to understand the concept but can be optimized as a single loop like this one :

arma::mat accumarray(arma::mat& subs, arma::vec& val)
{
    arma::u32 ar = arma::max(subs.col(0)) + 1;
    arma::u32 ac = arma::max(subs.col(1)) + 1;

    arma::mat A(ar, ac);
              A.zeros();

    for (arma::u32 r = 0; r < subs.n_rows; ++r)
        A(subs(r, 0), subs(r, 1)) += val(r);

    return A;
}

The only change are :

• init the output matrix with zero's.
• loop over subs rows to get the output indice(s)
• accumulate val to output (subs & val are row synchronized)

A 1-D version (vector) of the function can be something like :

arma::vec accumarray (arma::ivec& subs, arma::vec& val)
{
    arma::u32 num_elems = arma::max(subs) + 1;

    arma::vec A (num_elems);
              A.zeros();

    for (arma::u32 r = 0; r < subs.n_rows; ++r)
        A(subs(r)) += val(r);

    return A;
}

For testing 1D version :

arma::vec val = arma::regspace(101, 105);
arma::ivec subs;
subs << 0 << 2 << 3 << 2 << 3;
arma::vec A = accumarray(subs, val);
A.raw_print("A =");

The result is conform with matlab examples (see previous link)

A = 
    101
      0
    206
    208

This is not a strict copy of matlab accumarray function. For example, the matlab function allow to output vec/mat with size defined by sz that is larger than the intrinsec size of the subs/val duo.

Maybe that can be a idea for addition to the armadillo api. Allowing a single interface for differents dimensions & types.