退役后的分治练习

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

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

题意

告诉你(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] imes dp[i-1][(j-l) \% k] ]

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

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

优化1

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

注意到上面有个卷积！那不妨这个卷积优化一下。
但是(k)也不大，就算优化成(O(nklog 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
", ans[0]);
    return 0;
}

欢迎留言指出问题！