L3ungj's blog

Using the multinomial theorem in competitive programming
By L3ungj, 3 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.

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

»
3 years ago, # |
  Vote: I like it +21 Vote: I do not like it

orz

→ Reply
»
3 years ago, # |
  Vote: I like it +19 Vote: I do not like it

hatsune miku real!!!!!

→ Reply
»
3 years ago, # |
  Vote: I like it +5 Vote: I do not like it

I don't think this is very useful, at least in the applications that you have shown. Both approaches consume $$$\mathcal{O}(n\cdot k)$$$ space and run in the same complexity, but the prefix sums approach is very straightforward and doesn't require any math work by hand.

→ Reply
  • »
    »
    3 years ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    Yes, I agree with you. I hope to see some applications of this method in some future CP math problems.

    → Reply
  • »
    »
    3 years ago, # ^ |
      Vote: I like it +1 Vote: I do not like it

    I think this is the approach of many people starting CP like me. A simpler example is to calculate $$$\sum_{i=l}^ra_i$$$, if one doesn't know about prefix sum they are likely to work that out with maths formula. Sometimes handworks help reducing space and time complexit, too.

    → Reply
»
3 years ago, # |
  Vote: I like it +3 Vote: I do not like it

Pratice problem using multinomial theorem

→ Reply
»
3 years ago, # |
  Vote: I like it +1 Vote: I do not like it

I find it easier without using multinomial sums (and it's faster too), but nevertheless interesting to see another approach.

→ Reply
  • »
    »
    3 years ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    Yes, very fun to see different methods can give the same result.

    → Reply
»
3 years ago, # |
← Rev. 2   Vote: I like it +14 Vote: I do not like it

Here is another implementation using prefix sums, and it seems to me quite easy to write:

vector<long long> sum(4, 0);
sum[0] = 1;
for (int x : a) {
    for (int i = 3; i > 0; --i) {
        sum[i] += sum[i - 1] * x;
    }
}
cout << sum[3] << "\n";
→ Reply
»
3 years ago, # |
  Vote: I like it 0 Vote: I do not like it

You're on the right path but with more practice in pattern recognition, you should be able to figure out the formula without even thinking about the multinomial theorem for small powers.

For large powers, a general formula using the sums of powers can be constructed by GEM rather than by hand.

→ Reply