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]
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.