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

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

QUPC 2018 I - Buffalo

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

問題へのリンク

問題概要

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

  • 容量が  X, Y の容器を用意して、以下の操作によって水の量  K を汲み取れる
    • 操作 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;
}