Fast double sum involving Kronecker symbol

Question

I have three lists $$A=\{\alpha_1...\alpha_N\}$$ $$B=\{\beta_1...\beta_N\}$$ $$M=\{M_1....M_N\}$$

In practice $\alpha_i$ and $\beta_i$ are complex numbers with modulus smaller than 1. The $M_i$s are integers in between -6 and 6.

I want to compute the following sum:

$$S=2\sum_i\sum_j\alpha_i\beta_j\delta_{M^i,M^j}$$

My best attempt is basically

ComputeSum[A_, B_,M_] := 
 Block[{kroneckerList},
  k= Outer[KroneckerDelta, M, M];
  ak= A*kroneckerList;
  Total[ak.B]
  ]

Which I believe does the job (but in any case what I want to compute is the sum above).
The problem is that $N$ can be quite big, up to $3\times 10^6$. and the k= Outer[KroneckerDelta, M, M] line gets computationally expensive.

Do you have faster solutions ?

Answer 1

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, A}], First -> Last, Total]],
   Values[GroupBy[Transpose[{M, B}], First -> Last, Total]]
   ];


n = 10000;
A = RandomReal[{-1, 1}, n];
B = RandomReal[{-1, 1}, n];
k = 6;
M = RandomInteger[{-k, k}, n];

result = ComputeSum[A, B, M]; // AbsoluteTiming // First
result2 = ComputeSum2[A, B, M, k]; // 
  AbsoluteTiming // First
result3 = ComputeSum3[A, B, 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, A}], 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.