楽しいけど重くて、重いけど楽しい

問題概要

$N$ 要素から成る数列 $A_1, A_2, \dots, A_N$ が与えられる。 $N$ 個から 2 個選んでできる $C(N, 2)$ 通りのペア $X, Y$ のうち、以下の条件を満たすものを数え上げよ。

容量がの容器を用意して、以下の操作によって水の量を汲み取れる
- 操作 1: 一方の容器を水でいっぱいにする。
- 操作 2: 容器 $X$ からもう一方の容器 $Y$ に、 $Y$ がいっぱいになるか X が空になるまで水をうつす。
- 操作 3 : 一方の容器の中の水を全て捨てる。

制約

$2 \le N \le 3×10^{5}$
$2 \le K \le 2 × 10^{6}$
$1 \le A_i \le 10^{6}$

考えたこと

基本的には $X, Y$ の線形結合はだいたい作れるとみていいことから、

$K$ が ${\rm GCD}(X, Y)$ の倍数であること
$X + Y \ge K$ であること

が必要十分条件であることがわかる。例えば $K = 1$ だったら

全ペアを足して
素数 $p_1$ の倍数のペアを引いて
素数 $p_1, p_2$ に対して $p_1 p_2$ の倍数のペアを足して
...

という感じの約数系包除原理になる。 $K$ が一般でもあまり変わらず、 $K$ に含まれる素因子に対してはその指数を 1 個増やしたやつを足し引きすればいい。

最後に、

整数 $d$ の倍数のペアの個数 (のうち和が $K$ 以上であるもの)

を求める部分は、バケット法からのしゃくとり法 (ここでは二分探索でやった) などでできる。 $d$ の倍数の個数は最悪でも $\frac{A}{d}$ 個であることから、調和級数の和の計算量になる。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

const int MAX = 1000010;
bool IsPrime[MAX];
int MinFactor[MAX];
vector<int> preprocess(int n = MAX) {
    vector<int> res;
    for (int i = 0; i < n; ++i) IsPrime[i] = true, MinFactor[i] = -1;
    IsPrime[0] = false; IsPrime[1] = false; MinFactor[0] = 0; MinFactor[1] = 1;
    for (int i = 2; i < n; ++i) {
        if (IsPrime[i]) {
            MinFactor[i] = i;
            res.push_back(i);
            for (int j = i*2; j < n; j += i) {
                IsPrime[j] = false;
                if (MinFactor[j] == -1) MinFactor[j] = i;
            }
        }
    }
    return res;
}

vector<pair<long long,long long> > prime_factor(long long n) {
    vector<pair<long long,long long> > res;
    while (n != 1) {
        int prime = MinFactor[n];
        int exp = 0;
        while (MinFactor[n] == prime) {
            ++exp;
            n /= prime;
        }
        res.push_back(make_pair(prime, exp));
    }
    return res;
}


int GCD(int a, int b) { return b ? GCD(b, a%b) : a;}

int main() {
    int N;
    long long K;
    cin >> N >> K;
    vector<long long > A(MAX, 0);
    for (int i = 0; i < N; ++i) {
        long long  a; cin >> a;
        A[a]++;
    }
    
    preprocess();
    long long res = 0;
    for (long long d = 1; d < MAX; ++d) {
        auto pf = prime_factor(d);
        long long mebius = 1;
        for (int i = 0; i < pf.size(); ++i) {
            if (pf[i].second > 1) mebius = 0;
            else mebius *= -1;
        }
        if (!mebius) continue;

        long long nd = 1;
        for (int i = 0; i < pf.size(); ++i) {
            long long nk = K;
            int num = 0;
            long long p = pf[i].first;
            while (nk % p == 0) ++num, nk /= p, nd *= p;
            nd *= p;
        }
        vector<long long> vals;
        for (long long dd = nd; dd < MAX; dd += nd) {
            for (int i = 0; i < A[dd]; ++i) vals.push_back(dd);
        }
        long long num = vals.size();
        long long tmp = 0;
        for (int left = 0; left < vals.size(); ++left) {
            int right = lower_bound(vals.begin(), vals.end(), K - vals[left]) - vals.begin();
            right = max(right, left + 1);
            tmp += num - right;
        }
        
        res += tmp * mebius;
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

QUPC 2018 I - Buffalo

問題概要

制約

考えたこと