Return to Revisions

ComputeSum[A_, B_, M_] := A . Outer[KroneckerDelta, M, M] . B;

(*A helper function that for additive assembly of `SparseArray`s (_Mathematica's_ default is first in, last out.) *)
Options[MySparseArray] = {"Background" -> 0.};
MySparseArray[X_, r_, f_ : Total] := 
  If[(Head[X] === Rule) && (X[[1]] === {}),
   X[[2]],
   With[{spopt = SystemOptions["SparseArrayOptions"]},
    Internal`WithLocalSettings[
     SetSystemOptions[
      "SparseArrayOptions" -> {"TreatRepeatedEntries" -> f}],
     SparseArray[X, r, OptionValue["Background"]],
     SetSystemOptions[spopt]]
    ]
   ];

ComputeSum2[A_, B_, M_, k_] := Dot[
   MySparseArray[Partition[M + k + 1, 1] -> A, {2 k + 1}],
   MySparseArray[Partition[M + k + 1, 1] -> B, {2 k + 1}]
   ];

ComputeSum3[A_, B_, M_] := Dot[
   Values[GroupBy[Transpose[{M, \[Alpha]}], First -> Last, Total]],
   Values[GroupBy[Transpose[{M, \[Beta]}], First -> Last, Total]]
   ];


n = 10000;
\[Alpha] = RandomReal[{-1, 1}, n];
\[Beta] = RandomReal[{-1, 1}, n];
k = 6;
M = RandomInteger[{-k, k}, n];

result = ComputeSum[\[Alpha], \[Beta], M]; // AbsoluteTiming // First
result2 = ComputeSum2[\[Alpha], \[Beta], M, k]; // 
  AbsoluteTiming // First
result3 = ComputeSum3[\[Alpha], \[Beta], M]; // AbsoluteTiming // First
Abs[result - result2]
Abs[result - result3]
Abs[result - result3]

16.8646

0.002846

0.006937

2.27374*10^-12

9.09495*10^-13

Edit

The idea of the two implementations is the same. We want to compute $$\begin{aligned} \sum_{i=1}^{n} \sum_{j=1}^n \alpha_i \, \delta_{M_i,M_j} \, \beta_j &= \sum_{i=1}^{n} \sum_{j=1}^n \sum_{k=-6}^6 \alpha_i \, \delta_{M_i,k} \, \delta_{k,M_j} \, \beta_j \\ &= \sum_{k=-6}^6 \left( \sum_{i=1}^{n}\alpha_i \, \delta_{M_i,k} \right) \, \left( \sum_{j=1}^n \delta_{k,M_j} \, \beta_j \right) \\ & = u^T v, \end{aligned}$$ where $$ u_k = \sum_{i=1}^n \alpha_i \, \delta_{M_i,k} \qquad v_k = \sum_{j=1}^n \beta_j \, \delta_{M_j,k}. $$ The naive summation costs O(n^2); but each of u and v can be computed in O((2\,k +1) \, n) time. So the new algorithm has complexity $$ O(2\,(2\,k +1) \,n + (2\,k +1)) = O(2\,(2\,k +1)\,n). $$ So if the range of k is much smaller than n, then we can save quite many flops this way.

Hence we may use

MySparseArray[Partition[M + k + 1, 1] -> A, {2 k + 1}]

(where we have to add shift the integers in M to be all greater than 0) or

Values[GroupBy[Transpose[{M, \[Alpha]}], First -> Last, Total]]

to assemble the vector u. Likewise we can do it for v. And in the end we just have Dot u and v together to get the result.