L3ungj's blog

Using the multinomial theorem in competitive programming
By L3ungj, 2 years ago, In English

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 solution
Multinomial theorem solution

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

Implementation

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.

Full text and comments »

  • Vote: I like it
  • +81
  • Vote: I do not like it