退役后的分治练习

朋友和我说了一道美团面试题。让我对自己的智商产生了一些信心。

这是一篇不成熟的文章，请告诉我怎么减少摘要的最少字数。

题意

告诉你\(n,l,r,k\)，计算有多少个长度为\(n\)的数列，满足\(l \leq a\_i \leq r,i \in [1,n]\)，且\(k | \sum_i a_i\)。结果对\(1e9+7\)取模。

数据范围

\[n \leq 1e5,k \leq 100, l \leq r \leq 1e9 \]

基本思路

一个很好想到的做法是滚动的二维dp，\(dp[i][j]\)表示前\(i\)个数字对\(k\)模数为\(j\)的方案数。那么转移方程为

\[dp[i][j] = \sum dp[i-1][l]\times dp[i-1][(j-l) \% k] \]

\[dp[i] = dp[i-1] \otimes dp[i-1] \]

其中\(\otimes\)表示循环卷积。

优化1

上面算法复杂度为\(O(nk^2)\)，这个数据范围有点卡呢~

注意到上面有个卷积！那不妨这个卷积优化一下。
但是\(k\)也不大，就算优化成\(O(nk\log k)\)好像还是有点费劲。

而且还要带精度的FFT，如果模数是\(998244353\)我会毫不犹豫NTT。

优化2

只能换做法了？

这个题目其实每个位置不具有什么特殊性质，所以我们可以考虑把\(n\)拆成两段，后面一端是前面的延伸。这样看似乎有了新的方向！

我们分治整个序列。那么这个式子大概可以表示成这个样子。

\[f(n)=f(mid) \otimes f(n - mid) \]

这里卷积不做优化了，直接最暴力的上，于是我们得到了下面这个\(O(k^2 \log_2 n)\)的做法。

#include <bits/stdc++.h>
using namespace std;
const int N = 105;
typedef long long ll;
const ll mod = 1e9 + 7;

vector<ll> ans;

inline ll calcIn(ll r, int k, int i) {
    if (r < i)
        return 0;
    return (r - i) / k + 1;
}

vector<ll> tmp, one;

vector<ll> bake;
inline void conv(vector<ll>& a, vector<ll>& b, int k) {
    assert((int)a.size() == k);
    assert((int)b.size() == k);
    bake.clear();
    for (int i = 0; i < k; i++)
        bake.push_back(a[i]), a[i] = 0;

    for (int i = 0; i < k; i++)
        for (int j = 0; j < k; j++)
            a[(i + j) % k] = (a[(i + j) % k] + bake[i] * b[j] % mod) % mod;
}

void solve(int n, int k, ll l, ll r) {
    if (n == 1) {
        for (int i = 0; i < k; i++)
            one[i] = ans[i] = calcIn(r, k, i) - calcIn(l - 1, k, i);
        return;
    }
    int mid = n >> 1;
    solve(mid, k, l, r);
    tmp.clear();
    for (vector<ll>::iterator it = ans.begin(); it != ans.end(); it++)
        tmp.push_back(*it);
    assert(mid == n - mid || mid == n - mid - 1);
    if (mid == n - mid - 1)
        conv(tmp, one, k);
    conv(ans, tmp, k);
}

int main() {
    int n, k;
    ll l, r;
    scanf("%d%lld%lld%d", &n, &l, &r, &k);
    ans.clear();
    one.clear();
    ans.resize(k);
    one.resize(k);
    solve(n, k, l, r);
    printf("%lld\n", ans[0]);
    return 0;
}

欢迎留言指出问题！