Codeforces Round #428 Div.2 D - Winter is here

http://codeforces.com/contest/839/problem/D

問題概要

N要素の数列Aが与えられる。Aの部分集合Sに対して、strength(S) = |S|・gcd(S) と定義する。gcd(S) > 1 なるすべての部分集合Sに対するstrengthの総和を求めよ。

N <= 2 * 105, max(A) <= 106

解法

一旦問題を緩和して、gcdではなく約数で考えてみる。すべての要素がxを約数に持つような部分集合の中で最も要素数が多いものをS_xと置くと、緩和版のstrengthは

  f(x) = \sum_{T \subseteq S_x}{|T|\cdot x}

と表せる。

ここでS_xは「gcdがxである集合」の上位集合であるから、f(x)から「gcdがxでない集合Tの|T|・x」を引いてやれば本当に求めたかったgcdの場合での値を得ることができる。では「S_xの中でgcdがxでない集合」のgcdがどのようになるのかと考えてみると、それらは必ずxの倍数になる。ゆえにxが大きいものから順にf(x)を求めていき、そこからΣf(xの倍数)を引いていけば答えが求まる。

肝心のf(x)の値だが、部分集合の取り方から以下のような組合せの式が立てられる。

 f(x) = x \cdot ({}_{|S_x|} C_1 + { }_{|S_x|} C_{2} + \ldots + { }_{|S_x|} C_{S_x})

後半の組合せの和は二項係数の性質から 2^{|S_x|-1}と変形できる(らしい)ので、2の累乗を前計算しておくか繰り返し二乗法ですぐ求まる。

自分は以下のように実装した。

  1. Aの全要素について約数を調べ、どの数が何回約数として使われているかをカウントする
  2. 大きい方からf(x)を求め、 g(x) = f(x) - \sum_{i>=2}{f(ix) / i}を求める(割り算のためにmod109+7での逆元をあらかじめ求めておく)
  3. g(x)の合計をとる

計算量はO(N*sqrt(N) + max(A) * log(N))

感想

C++で書いてTL3000msのところ2970msとかだったので、想定解はもうちょっとスマートになるっぽい? →今見てきたら最速で100ms強でした。速すぎる…。ちなみに私が堂々の最下位です。本当にありがとうございました。

他の方の解説見ると逆元持ち出してるところは明らかに冗長で、個数だけで包除やったあと最後に掛け算すればよい、ということはわかった。あと約数カウントはエラトステネスっぽくやると速くできるっぽい。これが一番でかそう。

コード (C++)

#include <bits/stdc++.h>
using namespace std;
#define REP(i,n) for (int i=0;i<(n);i++)
#define REP2(i,m,n) for (int i=m;i<(n);i++)
typedef long long ll;

const ll MAX = 1000000;
const ll MOD = 1000000007;
ll A[202020];
ll cnt[MAX + 1];
ll pow2[MAX + 1];
ll modinv[MAX + 1];
ll ans[MAX + 1];


int main() {
    pow2[0] = 1;
    REP(i, MAX) pow2[i + 1] = pow2[i] * 2 % MOD;
    modinv[0] = modinv[1] = 1;
    REP2(i, 2, MAX + 1) modinv[i] = modinv[MOD % i] * (MOD - MOD / i) % MOD;

    int N; cin >> N;
    REP(i, N) cin >> A[i];
    memset(cnt, 0, sizeof(cnt));
    memset(ans, 0, sizeof(ans));

    REP(i, N) {
        for (ll j = 1; j * j <= A[i]; ++j) {
            if (j * j == A[i]) {
                cnt[j] += 1;
            } else if (A[i] % j == 0) {
                cnt[j] += 1;
                cnt[A[i] / j] += 1;
            }
        }
    }

    for (ll i = MAX; i > 1; i--) {
        if (!cnt[i]) continue;
        ans[i] += i * cnt[i] % MOD * pow2[cnt[i]-1] % MOD;
        ans[i] %= MOD;
        int hoge = 2;
        for (ll j = i + i; j <= MAX; j += i) {
            ans[i] = (ans[i] - ans[j] * modinv[hoge]) % MOD;
            ans[i] = (ans[i] + MOD) % MOD;
            hoge += 1;
        }
    }

    ll tmp = 0;
    REP(i, MAX+1) tmp = (tmp + ans[i]) % MOD;
    cout << tmp << endl;
}