Tokitsukaze and Three Integers

原题链接

TAG：容斥，动态规划

思路

\(dp[x]\) 表示 \(a[i]\times a[j]\equiv x\pmod{p}(i\neq j)\) 的数量，可以在 \(O(n^2)\) 的时间复杂度内算出。

接下来枚举 \(k\)，计算 \(ans[x]\) 满足 \((a[i]\times a[j]+a[k])\equiv x\pmod{p}(i\neq j)\) 的数量，时间复杂度为 \(O(np)\)

但是会将 \(i=k\) 或 \(j=k\) 的情况计算到答案中，需要减去这种情况，枚举 i=k 的情况，减去 \(ans[x]\) 满足 \((a[i]\times a[j]+a[i])\equiv x\pmod{p}(i\neq j)\) 的情况。\(j=k\) 的时候同理。时间复杂度为 \(O(n^2)\)

时间复杂度为 \(O(n^2+np)\)

代码

#include <bits/stdc++.h>
using namespace std;
using LL = long long;

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int n, p;
    cin >> n >> p;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        a[i] %= p;
    }

    vector<int> dp(p);
    // 计算 a[i]*a[j]%p=x 的数量
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            if (i == j) continue;
            dp[(a[i] * a[j] % p)]++;
        }
    }

    // 计算 (a[i]*a[j]+a[k])%p=x 的数量
    vector<LL> ans(p);
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < p; j++) {
            ans[(j + a[i]) % p] += dp[j];
        }
    }

    // 减去 i==k || j==k 的情况
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            if (i == j) continue;
            ans[(a[i] * a[j] + a[i]) % p] -= 2;
        }
    }
    for (auto x : ans) cout << x << " ";

    return 0;
}