i_love_sqrt_decomp's blog

By i_love_sqrt_decomp, history, 9 months ago, In English

I just did a problem that asks you for $$$q$$$ queries of $$$[x^k](\prod_{i=1}^{n}(1+x^i))^2\ mod \ 10^9 + 7$$$. $$$n$$$ is given beforehand and $$$k$$$ is given in the queries. Constraints: $$$n,k,q \leq 10^5$$$. Is there a simple solution that doesn't involve a huge polynomials template that runs in $$$o(NK)$$$? (The time limit is 3 seconds).

My approach
  • Vote: I like it
  • +19
  • Vote: I do not like it

»
6 months ago, hide # |
 
Vote: I like it +2 Vote: I do not like it

The original problem was about choosing $$$2$$$ subsets of $$$\{1,2,...,n\}$$$ such that the sum of the $$$2$$$ subsets is $$$k$$$. The problem can be solved in $$$O(k\sqrt{k})$$$ time using DP, with $$$DP_{i,j}$$$ being the number of ways such that the sum is $$$j$$$ and the total size of the two sets are $$$i$$$.

We have this DP formula: $$$DP_{i,j} = DP_{i,j-i} +2\cdot DP_{i-1,j-i} +DP_{i-2,j-i}$$$, which can be interpreted as 4 cases:

  • Add $$$1$$$ to every element in both sets.
  • Add a $$$0$$$ to the 1st set, then add $$$1$$$ to every element in both sets.
  • Add a $$$0$$$ to the 2nd set, then add $$$1$$$ to every element in both sets.
  • Add a $$$0$$$ to both sets, then add $$$1$$$ to every element in both sets.

But the numbers in the sets can exceed $$$N$$$, so we need to consider 3 cases:

  • The 1st set have a $$$N+1$$$ in it.
  • The 2nd set have a $$$N+1$$$ in it.
  • Both sets have a $$$N+1$$$ in it.

Our finalized formula, which consider those 3 cases, is $$$DP_{i,j} = DP_{i,j-i} +2\cdot DP_{i-1,j-i} + DP_{i-2,j-i} - 2 * DP_{i-1, j-n-1} - DP_{i-2, j - 2\cdot (n+1)}$$$.

Since $$$i$$$ is $$$O(\sqrt{k})$$$ (consider the number of elements that we can choose from $$$\{1,1,2,2,...,n,n\}$$$ so that the sum does not exceed $$$k$$$), we got an $$$k\sqrt{k}$$$ solution.

»
6 months ago, hide # |
 
Vote: I like it -6 Vote: I do not like it

Just like segment tree.
If $$$\displaystyle\prod_{i=1}^{\frac{n}{2}}f_i(x),\displaystyle\prod_{i=\frac{n}{2}+1}^nf_i(x)$$$ have been calculated,the complexity of timing them is $$$O(n\log n)$$$.
As the Master Theorem,the time complexity of the whole algorithm is $$$T(n)$$$,merging is $$$O(n\log n)$$$,so $$$T(n)=2T(\frac{n}{2})+O(n\log n)=O(n\log^2n)$$$.
Noticed that $$$k\le 10^5$$$ ,so you only need to calculate the answer $$$\bmod x^{10^5+1}$$$.
After doing this, you will get the polynomial $$$\prod_{i=1}^n (1+x^i))$$$.
Then get the square of it, and get the answer.
You have $$$q$$$ querys, so actually it's $$$O(n\log^2 n+q)$$$.

  • »
    »
    6 months ago, hide # ^ |
     
    Vote: I like it 0 Vote: I do not like it

    That is much better than $$$O(k\sqrt{k})$$$.

    But noticed that the modulo is $$$10^9+7$$$.
    You should use 3 NTTs + CRT or MTT

  • »
    »
    6 months ago, hide # ^ |
     
    Vote: I like it 0 Vote: I do not like it

    This does not work, the polynomials you get in the recursion are (much) larger than $$$O\left(\frac{n}{2}\right)$$$. This would be at least $$$O(n^2)$$$, with extra log factor from FFT.

»
6 months ago, hide # |
Rev. 2  
Vote: I like it +6 Vote: I do not like it

There is a trick to do this in $$$O(k log(k))$$$, where $$$k$$$ is the largest value in the queries ($$$\leq 10^5$$$). It goes as follows:

$$$\left(\prod_{i=1}^{n}(1+x^i)\right)^2 = e^{\log\left(\prod_{i=1}^{n}(1+x^i)^2\right)} = e^{2\sum_{i=1}^{n}\log\left(1+x^i\right)}.$$$

Now we use that $$$\log(1+x) = x - \frac{x^2}{2} + \frac{x^3}{3} - \frac{x^4}{4} + \ldots$$$. To get $$$\log(1+x^i)$$$, we replace $$$x$$$ in this formula by $$$x^i$$$. Since we are only interested in the first $$$k \leq 10^5$$$ coefficients, the number of nonzero coefficients of $$$\log(1+x^i)$$$ that we care about is $$$O\left(\frac{k}{i}\right)$$$. Therefore, we can calculate the first $$$k$$$ coefficients of $$$2\sum_{i=1}^{n}\log\left(1+x^i\right)$$$ in $$$O\left(\sum_{i=1}^{n} \frac{k}{i}\right) = O(k \log(n))$$$. Lastly, once we calculated $$$p(x) = 2\sum_{i=1}^{n}\log\left(1+x^i\right)$$$, we need to calculate the generating function $$$e^p(x)$$$. This can be done in $$$O(k \log(k))$$$ using FFT (link).

Since any terms $$$(1+x^i)$$$ with $$$i \gt k$$$ are irrelevant for the answer, we can assume that the complexity is therefore $$$O(k\log(k))$$$, where $$$k$$$ is the largest index we care about. The only annoying thing is that $$$10^9+7$$$ is not suited for NTT, but you can still do it, only with a ever so slightly worse constant factor.