I recently encountered a problem that requires to compute $$$\sum_{i=0}^{n-1} \sum_{j=i+1}^{n-1} \sum_{k=j+1}^{n-1} a_i a_j a_k$$$ in $$$O(n)$$$. My friend __declspec, (orz), uses prefix sums to compute it, which takes many time to implement, and also easy to have bugs (to me). I found the product of the 3 elements familiar in the multinomial theorem, so I tried applying it to the problem.
The multinomial theorem
From Wikipedia, the theorem states
$$$ (\sum\limits_i {a_i})^n = \sum\limits_{K=n} \frac{n!}{\prod\limits_i{k_i!}} \prod\limits_i {a_i}^{k_i} $$$where
$$$ K = \sum\limits_i k_i $$$ and $$$ k_i \ge 0 $$$ for all $$$i$$$Example
In the case where $$$n=2$$$,
$$$ S^2 = \sum\limits_i {a_i^2} + 2\sum\limits_i \sum\limits_{j>i} a_i a_j$$$where
$$$S = \sum\limits_i {a_i}$$$In the recent problem 1637D - Yet Another Minimization Problem, $$$\sum\limits_i \sum\limits_{j>i} {(a_i + a_j)^2}$$$ shows up. According to the editorial (thanks for the amazing problems and editorial), it can be simplified as follow,
$$$\sum\limits_i \sum\limits_{j>i} {(a_i + a_j)^2} = \sum\limits_i \sum\limits_{j>i} {(a_i^2 + 2a_i a_j + a_j^2)}$$$The middle term can be simplified,
$$$= \sum\limits_i (n-i-1)a_i^2 + \sum\limits_i \sum\limits_{j>i} a_j^2 + S^2 - \sum\limits_i {a_i^2} $$$The middle term can be simplified with some change of indices,
$$$ \sum\limits_i \sum\limits_{j>i} a_j^2 = \sum\limits_j \sum\limits_{i<j} a_j^2 = \sum\limits_j ja_j^2$$$Thus,
$$$ \sum\limits_i \sum\limits_{j>i} {(a_i + a_j)^2} = (n-2)\sum\limits_i a_i^2 + S^2.$$$The main dish — triple summation
First we let $$$S = \sum\limits_i {a_i}$$$. By the multinomial theorem,
$$$ S^3 = \sum\limits_i a_i^3 + 3 \sum\limits_i a_i^2 \sum\limits_{j \neq i} a_j + 6\sum\limits_{i} \sum\limits_{j>i} \sum\limits_{k>j} a_i a_j a_k $$$ $$$ = \sum\limits_i a_i^3 + 3 \sum\limits_i a_i^2 (S - a_i) + 6\sum\limits_{i} \sum\limits_{j>i} \sum\limits_{k>j} a_i a_j a_k $$$ $$$ = 3S \sum\limits_i a_i^2 -2 \sum\limits_i a_i^3 + 6\sum\limits_{i} \sum\limits_{j>i} \sum\limits_{k>j} a_i a_j a_k$$$Rearranging,
$$$ \sum\limits_{i} \sum\limits_{j>i} \sum\limits_{k>j} a_i a_j a_k = \frac{1}{6}(S^3 + 2 \sum\limits_i a_i^3 -3S \sum\limits_i a_i^2 ) $$$Comparison to prefix sums
Prefix sum solutionfor (int i = 0; i < n; ++i)
p[i + 1] = p[i] + a[i];
for (int i = 0; i < n; ++i)
p2[i + 1] = p2[i] + a[i] * p[i];
int ans = 0;
for (int i = 0; i < n; ++i)
ans += a[i] * p2[i];
cout << ans << '\n';
Multinomial theorem solutionint s = 0;
for (int i = 0; i < n; ++i)
s += a[i];
int ans2 = s * s * s;
for (int i = 0; i < n; ++i)
ans2 += 2 * a[i] * a[i] * a[i] - 3 * a[i] * a[i] * s;
cout << ans2 / 6 << '\n';
Which do you think is better? Comment your thoughts below!
Why stop at n=3?
By the multinomial theorem,
$$$(\sum\limits_i a_i)^4 = \sum\limits_i a_i^4 + 4\sum\limits_i a_i^3 \sum\limits_{j \neq i} a_j + 6\sum\limits_{a_i^2} \sum\limits_{j>i} {a_j^2} + 12 \sum\limits_i a_i^2 \sum\limits_{j \neq i} a_j \sum\limits_{k > j \wedge k \neq i} a_k + 24 \sum\limits_i \sum\limits_{j>i} \sum\limits_{k>j} \sum\limits_{l>k} a_i a_j a_k a_l$$$At this point we can define $$$S_j = \sum\limits_i {a_i^j}$$$. Note that:
$$$4\sum\limits_i a_i^3 \sum\limits_{j \neq i} a_j = 4\sum\limits_i a_i^3 (S_1 - a_i) = 4S_3 S_1 - 4S_4$$$ $$$6\sum\limits_{a_i^2} \sum\limits_{j>i} {a_j^2}= 6 \cdot \frac{1}{2}(S_2^2 - S_4) = 3S_2^2 - 3S_4$$$ $$$12 \sum\limits_i a_i^2 \sum\limits_{j \neq i} a_j \sum\limits_{k > j \wedge k \neq i} a_k = 12\sum\limits_i a_i^2 (\sum\limits_{j} \sum\limits_{k>j} a_j a_k - a_i \sum\limits_{j \neq i} a_j) $$$ $$$= 12\sum\limits_i a_i^2 (\frac{1}{2} (S_1^2 - S_2) - a_i (S_1 - a_i)) = 6S_1^2 S_2 - 6S_2^2 -12S_3 S_1 + 12 S_4$$$Putting them altogether,
$$$ S_1^4 = 6S_4 - 8S_3 S_1 - 3S_2^2 + 6S_2 S_1^2+ 24 \sum\limits_i \sum\limits_{j>i} \sum\limits_{k>j} \sum\limits_{l>k} a_i a_j a_k a_l$$$Thus,
$$$\sum\limits_i \sum\limits_{j>i} \sum\limits_{k>j} \sum\limits_{l>k} a_i a_j a_k a_l = \frac{1}{24} (S_1^4 - 6S_4 + 8S_3 S_1 + 3S_2^2 - 6S_2 S_1^2) $$$Well, using prefix sums might be better instead of going through these tedious calculations...
Implementationlong long ans2 = 0;
vector<int> s(4);
for (int j = 0; j < 4; ++j)
for (int i = 0; i < n; ++i)
s[j] += pw(a[i], j + 1);
ans2 = pw(s[0], 4) - 6 * s[3] + 8 * s[0] * s[2] + 3 * pw(s[1], 2) - 6 * pw(s[0], 2) * s[1];
cout << ans2 / 24 << '\n';
Conclusion
Using either prefix sums or multinomial theorem can compute the above summations in $$$O(n)$$$. Although the multinomial theorem solution requires tedious calculations, it can be used to simplify some summations in problems like the above example problem.
Thanks for reading and now let's take a rest to enjoy your daily dose of waifu. 
Полный текст и комментарии »
