AtCoder Grand Contest 006 C - Rabbit Exercise

http://agc006.contest.atcoder.jp/tasks/agc006_c

問題概要

1~Nの番号が振られたN匹のうさぎが一次元の数直線上に存在し、各うさぎの初期座標を示す数列Xが与えられる。うさぎi(2<=i<=N-1)がジャンプを行うと、うさぎ(i-1)かうさぎ(i+1)のどちらかが等確率で選ばれ、うさぎiは選ばれた方のうさぎに関して対称な座標に移動する。ここでうさぎの番号を要素にもつM要素の数列Aが与えられ、うさぎたちがこの順番にジャンプすることを1セットとする。ジャンプをKセット行うとき、各うさぎの最終的な座標の期待値を求めよ。

3 <= N <= 105

1 <= M <= 105

1 <= K <= 1018

解法

うさぎaが(a+1)の方にジャンプするとジャンプ後の座標は2X[a+1] - X[a]になり、(a-1)の方にジャンプすると2X[a+1] - X[a]になるから、1回のジャンプ後のaの座標の期待値は2つの平均をとってX[a+1] + X[a-1] - X[a] となる。この計算を繰り返していけば(実行時間はともかく)答えは求まる。あとは高速化が問題で、Kの大きさから行列累乗をしたくなるがO(M3logK)とかになってしまって無理。

ここでの正解は階差数列で考えることである。Xの階差数列Yを取ると、ジャンプの操作は以下のようになる。

  • X: X[i-1], X[i], X[i+1] → X[i-1], X[i-1] - X[i] + X[i+1], X[i+1]
  • Y: X[i] - X[i-1], X[i+1] - X[i] → X[i+1] - X[i], X[i] - X[i-1]

これをみると階差数列の方では前後の値が単にswapされていることがわかる。つまりジャンプ操作は階差数列の2要素のswapであるとみなせるようになる。このswapを1セット分シミュレートすれば、1セットでどの要素がどこに移動するかがわかる。これは要するに置換なので、あとはK回分の置換の積を繰り返し2乗法の要領で求めていけばO(MlogK)ですべてのジャンプが完了したあとの階差数列を得ることができる。ここで元の数列において先頭のうさぎの座標が変わることはないので、最後は階差数列の要素を順に足していくだけで元の数列を復元することができる。

感想

難しい問題あるある、1個だけあります

♪~~~~~♪~~~~~~~

階差とりがち

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
long[] perm_prod(const long[] A, const long[] P) {
    auto ret = A.dup;
    foreach (i; 0..A.length) {
        ret[i] = A[P[i].to!int];
    }
    return ret;
}
 
long[] perm_pow(const long[] A, long x) {
    long[] P = A.dup;
    long[] ret = new long[](P.length);
    foreach (i; 0..P.length) ret[i] = i;
    while (x) {
        if (x % 2) ret = perm_prod(ret, P);
        P = perm_prod(P, P);
        x /= 2;
    }
    return ret;
}
 
void main() {
    auto N = readln.chomp.to!int;
    auto X = readln.split.map!(to!long).array;
    auto s = readln.split.map!(to!long);
    auto M = s[0].to!int;
    auto K = s[1];
    auto A = readln.split.map!(x => x.to!int - 1).array;
 
    auto D = new long[](N - 1);
    foreach (i; 0..N-1) D[i] = X[i + 1] - X[i];
 
    auto B = new long[](N - 1);
    foreach (i; 0..N-1) B[i] = i;
    foreach (i; 0..M) swap(B[A[i] - 1], B[A[i]]);
 
    B = perm_pow(B, K);
    D = perm_prod(D, B);
 
    auto ans = new long[](N);
    ans[0] = X[0];
    foreach (i; 0..N-1) ans[i + 1] = ans[i] + D[i];
 
    ans.each!writeln;
}

AtCoder Beginner Contest 020 D - LCM Rush

http://abc020.contest.atcoder.jp/tasks/abc020_d

問題概要

整数N, Kが与えられる。  \sum_{i=1}^{N}{LCM(i, K)}を求めよ。

1 <= N, K <= 109

解法

(要点:「約数にxをもつ集合」から「約数に2xを持つ集合」, 「約数に3xを持つ集合」, …を引いていけば「GCDがxの集合」を出せる)

LCM(n, K) = nK / GCD(n, K) である。ここでGCDの定義より、GCD(n, K)が取る値は必ずKの約数のうちのどれかになる。Kの約数ごとに場合分けして求める値を計算することにすると、K / GCD(n, K)が定数になるためΣの外にくくりだすことができる。以上よりこの問題は結局、Kのすべての約数xについて「KとのGCDがxであるような整数i (1 <= i <= N)の和 × K × x」を求め、それらを足し合わせる問題と言い換えることができる。なお109以下の数であれば約数の個数は最大でも1000ちょっとくらいしかないので、これは十分に列挙可能である。

残る問題は「KとのGCDがxであるような整数i (1 <= i <= N)の和」をどう計算するかということである。ここで一旦GCDという条件を緩めて「約数にxを持つような整数i (1 <= i <= Nの集合」がどうなるかを考えてみると、単純な等差数列になる。例えばx = 1のときは1, 2, …, Nになるし、x = 2のときは2, 4, …となる(形式的には初項x, 公差x, 項数[N/x]の等差数列)。ここで「約数にxをもつ集合」から「約数に2xを持つ集合」「約数に3xを持つ集合」…を引いていくと、残ったものが「GCDがxである集合」になる(ただし2x, 3x等がKの約数でない場合は飛ばす)。であるから、Kの約数xの大きい方から見ていって「約数にxを持つ集合」(=等差数列)の和を求め、Kの約数であるような2x, 3x, …, の場合の値を引いていけばKのすべての約数xについてGCD(i, K) = xの場合の値が出せる。最後にそれらを足し合わせたものが答え。

等差数列の和は公式に当てはめればO(1)で、Kの約数の個数は高々1000個くらいになるので十分に間に合う。

感想

前のこどふぉで解いた問題が「LCM Rushっぽい」と言われていたのでそろそろ解けるかなと思ってやってみたら、解けた。これでABC-Dが全埋めということになって嬉しい。

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
immutable long MOD = 10^^9 + 7;
 
void main() {
    auto s = readln.split.map!(to!long);
    auto N = s[0];
    auto K = s[1];
 
    long[] F;
    for (long i = 1; i * i <= K; i++) {
        if (i * i == K) F ~= i;
        else if (K % i == 0) F ~= i, F ~= K / i;
    }
    F.sort();
 
 
    auto ans = new long[](F.length);
    auto sums = new long[](F.length);
    
    for (int i = F.length.to!int - 1; i >= 0; i--) {
        long n = N / F[i];
        sums[i] = n * (2 * F[i] % MOD + (n - 1) * F[i] % MOD) % MOD * powmod(2, MOD-2, MOD) % MOD;
 
        for (int j = i + 1; j < F.length; j++) {
            if (F[j] % F[i] == 0) {
                sums[i] -= sums[j] % MOD;
                sums[i] = (sums[i] + MOD) % MOD;
            }
        }
        
        ans[i] = K * sums[i] % MOD * powmod(F[i], MOD-2, MOD) % MOD;
    }
 
    long ansans = 0;
    foreach (a; ans) ansans = (ansans + a) % MOD;
    ansans = (ansans + MOD) % MOD;
    ansans.writeln;
}
 
 
long powmod(long a, long x, long m) {
    long ret = 1;
    while (x) {
        if (x % 2) ret = ret * a % m;
        a = a * a % m;
        x /= 2;
    }
    return ret;
}

Codeforces Round #429 Div.2 D / Div.1 B - Leha and another game about graph

http://codeforces.com/contest/840/problem/B

問題概要

N頂点M辺の連結な無向グラフGと、N要素の数列Dが与えられる。数列Dの要素は-1, 0, 1のいずれかである。Gから任意の辺集合を選び、選ばれた辺以外を削除するとき、以下の条件を満たすことができるか判定せよ。またできるのであればそのような辺の集合を構成せよ。

  • D[i]が0なら頂点iの次数は偶数
  • D[i]が1なら頂点iの次数は奇数
  • D[i]が-1なら頂点iの次数はなんでもよい

N, M <= 3 * 105

解法

(次数はすべてmod2で述べる。)

まずグラフの次数の合計は必ず偶数になる(辺1本につき合計次数が2増えるから)ため、Dに-1がなくかつDの合計が奇数の場合は条件は満たせない。

一度全部の辺を削除して、辺を張り直す問題として考えてみる。ここでひとつのパスを選んでそのパスに含まれる辺を張ることにすると、パスの端点のみ次数が1になり、途中の点は0のままになる。このやり方で次数が1の点を2つずつ増やしていくことができる。もしパスを張る過程で以前張ったパスと同じ辺を通ることになったら、単にその辺を消せば次数は保存される。結局この問題は「D[i]=1なる点を2つとってきてそれらをつなぐパスを見つけることをD[i]=1の頂点がなくなるまで繰り返す」問題と同等になる。

しかし一般的な無向グラフ上で愚直にペアを作ってDFSなりBFSでそのパスを見つけ…という操作を行っていると最悪でO(N2)かかる(ペアの数*DFSの計算量)。ここでパスを簡単に構成できる形をしたグラフがあって、それは木である。グラフは連結であることが保証されているため、一部の辺だけ取ってくるとグラフ上で木を作ることができる。この木において「(D[i]+sum(dfs(iの子ども))) mod 2を返す、このとき返り値が0なら親との間の辺を有効にする」というDFSで解が構成できる。-1はつじつま合わせで適当に0か1のどちらかに寄せればいい。

感想

説明難しすぎて自分でも何言ってるかわからないんですが、1の間をパスでつなぐんだっていうのとパスの構成は木にすると簡単っていうのがたぶん勘所じゃないかと思います。

コード (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;

int N, M;
int D[300010];
vector<pair<int, int>> edges[300010];
int edge_states[300010];
bool used[300010];

void dfs1(int n) {
    used[n] = true;
    for (auto m: edges[n]) {
        if (used[m.first]) continue;
        edge_states[m.second] = 0;
        dfs1(m.first);
    }
}

int dfs2(int n, int p, int pe) {
    int cnt = D[n] == 1;
    for (auto m: edges[n]) {
        if (edge_states[m.second] == -1 || m.first == p) continue;
        cnt += dfs2(m.first, n, m.second);
    }
    if (cnt % 2 == 1) edge_states[pe] = 1;
    return cnt % 2;
}

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

    cin >> N >> M;
    REP(i, N) cin >> D[i];
    REP(i, M) {
        int a, b;
        cin >> a >> b;
        edges[a - 1].push_back(make_pair(b - 1, i));
        edges[b - 1].push_back(make_pair(a - 1, i));
    }

    int cnt1 = 0;
    int cnt_1 = 0;
    REP(i, N) if (D[i] == 1) cnt1 += 1; else if (D[i] == -1) cnt_1 += 1;
    if (cnt1 % 2 == 1 && cnt_1 == 0) {
        cout << -1 << endl;
        return 0;
    } else if (cnt1 % 2 == 1 && cnt_1 > 0) {
        REP(i, N) if (D[i] == -1) {D[i] = 1; break;}
    }

    REP(i, M) edge_states[i] = -1;
    memset(used, 0, sizeof(used));

    dfs1(0);
    dfs2(0, -1, 0);

    int ans = 0;
    REP(i, M) if (edge_states[i] == 1) ans += 1;

    if (ans == 0) {
        cout << ans << endl;
        return 0;
    }

    cout << ans << endl;
    REP(i, M) if (edge_states[i] == 1) cout << i + 1 << "\n";
}

AtCoder Regular Contest 081 F - Flip and Rectangles

http://arc081.contest.atcoder.jp/tasks/arc081_d

問題概要

H×Wの2次元のマス目が与えられる。各マスは白か黒のどちらかで塗られている。任意の行または列を選びその行または列に含まれるすべてのマスの色を反転させる、という操作を好きなだけ行えるとき、黒のマスだけで構成できる長方形の最大面積を求めよ。

H, W <= 2000

解法

なぜそうなるかは一旦置いて、解法のみ先に述べる。まず各行独立にxorで階差数列をとり、H×(W-1)の行列をつくる。この行列のj列目を見たときi行目~(i+x)行目の要素が全部同じであれば、元の行列のi行目~(i+x)行目, j列目~(j+1)列目を全部黒で塗りつぶして(x+1)×2の長方形を作ることができる。逆に異なる要素が出てきた場合、それ以上長方形を伸ばすことはできない。これは各列で独立に成り立つ。以上より解の構成は以下の手順で行うことができる。

  1. 各行独立にxorで階差数列をとってH×(W-1)の行列を作る
  2. 作った行列の各要素ごとに、「下方向にどれだけ同じ要素が続くか」を計算する
  3. 各行ごとに2で計算した値をヒストグラムのように見て「ヒストグラム中の長方形の最大面積」を計算する

2はO(HW)でできる。3も各行の面積最大化をO(W)で行えるのでO(HW)。あわせてO(HW)で答えが出る。コーナーケースとして、一行・一列だけで作れる最大の長方形を考慮に入れる必要がある(必ずH, Wになる)。以上で答えが求まった。

で、なぜそうなるのか?まず差分が同一なら長方形を縦に伸ばしていけることについて。差分が0であるのは横並びが白白または黒黒というパターンである。これが縦に並んでいるなら全部黒にできるのは自明である。逆に差分が1のとき、つまり横並びが白黒または黒白というパターンでは、一度どちらかの列に反転を入れればすべて白白または黒黒になる。あとは白白の行にだけ反転を入れていけば全部黒にできる。例↓

f:id:fluffyowl:20170822182857p:plain

逆に差分が異なる場合は長方形を伸ばすことが不可能である(公式解説でいう「悪い部分」)。まず横に反転操作を入れても階差行列の値に変化は起こらない。また縦に反転を入れると階差行列は列方向に全部同時に反転が起こる。ゆえにどうやっても上下の片方だけ差分を変える、ということができず、両方を差分0で揃えることができない=黒黒にできない、ということになる。以上が差分が同じである限り縦に長方形を伸ばしていける理由である。

もう一点、以上の議論が列ごとに独立であることについて。先ほども述べた通り横方向に反転操作を行っても階差行列の方には一切の変化が生じない。また縦に入れても1列すべてが同時に反転するので「縦に見て同じかどうか」という関係には変化が生じない。以上より長方形を作るための反転操作をどのように行っても長方形作成にかかわる条件には影響がない。よって列ごとに独立に見ることができる。

感想

前回のARC-Fで「0/1数列の区間flipを考えるときは階差をみてみる」という考え方を知ったばかりなのに解けなかった。本番中は2次元で階差ってどうやるの???ふえ〜~んで終了。それがわかったところで長方形の面積最大化を知らないし思いつけない。ムズい。とりあえず、ヒストグラムの長方形面積最大化を知れたのと、階差をとるとはつまり変化を考えてそれを形式的に表すことである、という認識は得た。今後に期待

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;

void main() {
    auto s = readln.split.map!(to!int);
    auto H = s[0];
    auto W = s[1];
    auto A = H.iota.map!(_ => readln.chomp).array;
    auto B = new int[][](H, W);
    auto C = new int[][](H, W-1);
    foreach (i; 0..H) foreach (j; 0..W-1) C[i][j] = A[i][j] == A[i][j+1];

    fill(B[H-1], 1); B[H-1][W-1] = 0;
    for (int j = 0; j < W - 1; j++) {
        for (int i = H - 2; i >= 0; i--) {
            if (C[i+1][j] == C[i][j]) B[i][j] = B[i+1][j] + 1;
            else B[i][j] = 1;
        }
    }

    
    int ans = max(H, W);
    
    foreach (i; 0..H) {
        DList!(Tuple!(int, int)) stack;
        foreach (j; 0..W) {
            if (stack.empty || stack.back[1] < B[i][j]) {
                stack.insertBack(tuple(j, B[i][j]));
            } else {
                int pos = j;
                while (!stack.empty && stack.back[1] > B[i][j]) {
                    int area = (j - stack.back[0] + 1) * stack.back[1];
                    ans = max(ans, area);
                    pos = stack.back[0];
                    stack.removeBack;
                }
                stack.insertBack(tuple(pos, B[i][j]));
            }
        }
    }

    ans.writeln;
}

AtCoder Regular Contest 081 E - Don't Be a Subsequence

http://arc081.contest.atcoder.jp/tasks/arc081_c

問題概要

文字列Sが与えられる。Sの部分列ではないような文字列のうち、最も短いものを求めよ。答えが複数ある場合には辞書順で最小のものを答えること。文字列はいずれも英小文字のみからなるものとする。

|S| <= 2 * 105

解法

Sに含まれていない文字がある場合にはその中で辞書順最小の1文字を答えればよい。もし26文字すべてがSに含まれているなら、少なくとも2文字は使う必要がある。ここで仮に1文字目に文字αを選択した場合、「文字αが最初に現れるインデックス+1」までジャンプしてまた再帰的に同じ問題を解くことになる。ゆえに「文字αがi文字目以降で最初に現れる場所」を前計算で求めておけば、あとはDPで解くことができる。

形式的には

next(i, α) = i文字目以降で文字αが初めて登場するインデックス

とすると、以下のDPで最短の文字数を求めることができる。


dp(i) = min(dp(next(i, \alpha) + 1)) + 1

(ただしSのi文字目以降に出てこない文字が存在するとき, dp(i) = 1)

具体的な文字列はdpのminを取るとき文字も一緒に記録しておけば最後に復元できる。

感想

最近たまたま解いていたTDPCの辞書順と同じようなやり方で出来た。部分文字列構築の典型手法と言える?

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;

void main() {
    auto S = readln.chomp;
    auto N = S.length.to!int;

    auto next = new int[][](26, N);
    foreach (i; 0..26) fill(next[i], -1);

    foreach (c; 0..26) {
        for (int i = N - 1; i >= 0; i--) {
            if (S[i] - 'a' == c) {
                next[c][i] = i + 1;
            } else if (i < N - 1) {
                next[c][i] = next[c][i+1];
            }
        }
    }

    auto mem = new Tuple!(int, int)[](N);
    fill(mem, tuple(1 << 29, 1 << 29));

    Tuple!(int, int) dfs(int n) {
        if (n == -1) return tuple(0, -1);
        if (n >= N) return tuple(1, -1);
        if (mem[n][0] != 1 << 29) return mem[n];
        auto ret = tuple(1 << 29, 1 << 29);
        foreach (i; 0..26) {
            auto t = dfs(next[i][n]);
            if (t[0] < ret[0]) ret = tuple(t[0], i);
        }
        mem[n] = tuple(ret[0] + 1, ret[1]);
        return mem[n];
    }

    dfs(0);
    string ans = "";

    for (int i = 0; i != -1;) {
        if (mem[i][1] == -1) break;
        char c = (mem[i][1].to!char + 'a').to!char;
        ans ~= c;
        i = next[c - 'a'][i];
    }

    ans.writeln;
}

CS Academy Round #42 (Div. 2 only) E - Xor Submatrix

https://csacademy.com/contest/round-42/task/xor-submatrix/

問題概要

N要素の数列UとM要素の数列Vが与えられる。これらを用いて,  A_{i, j} = V_i xor U_j となるようなN*Mの行列Aを構成する。Aから任意のサイズの部分行列を取り出し全要素のxorを取った時に得られる最大の値を求めよ。

1 <= N, M <= 1000, 0 <= V_i, U_i <= 229

解法

①部分行列の行または列が偶数のとき

同じもの2つで打ち消し合うというxorの性質上、部分行列を偶数行取った場合にはVの要素は関係なくなり、Uの部分列のxorを取っているのと同じことになる。偶数列取った場合も同じ。

②部分行列の行と列がともに奇数のとき

この場合も打消しを考えると結局Uの部分列(長さ奇数)とVの部分列(長さ奇数)のxorになる。

以上をまとめると求めたい値の候補は以下になる

  1. Uの偶数長の部分列のxor和
  2. Vの偶数長の部分列のxor和
  3. (Uの奇数長の部分列のxor和)xor(Vの奇数長の部分列のxor和)

このうち1, 2はO(N2)で列挙できる。3は単純に見ようとするとO(N4)かかってしまうので工夫が必要になる。

ここでxorを大きくするためにはどうすればよいかを考えると、上位の桁のbitがなるべく立つようにしていけばよいことがわかる。つまり集合Aと値xが与えられたとき、(Aの任意の1要素 xor x)の値を最大化するようなAの要素は、bitを上から見ていけば貪欲に選択できるということになる。

これを効率よく行うために、自分は左の子/右の子が各桁の0/1に対応するような2分木を作った。これである値がクエリとして与えられたとき、木に存在する値の中でxorを最大にするものがlog(集合の要素数)で取ってこれるようになる。この木に「Vの奇数長の部分列」によってできるすべてのxor和を詰めておいて、「Uの奇数長の部分列のxor和」をクエリとして投げれば上の箇条書きの3番の値もO(N2 logN)で全部見ることができる。

感想

2分木を自前のクラスで一生懸命作ってたらめちゃくちゃバグらせてしまった。vectorで十分だったかもしれない

コード (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;

class BinaryNode {
public:
    int data = -1;
    BinaryNode* left = nullptr;
    BinaryNode* right = nullptr;
};


int N, M;
int A[2020];
int B[2020];

BinaryNode* construct(vector<int> p, int k) {
    auto bn = new BinaryNode;
    if (p.size() == 1) {
        bn->data = p[0];
        return bn;
    }
    vector<int> l = vector<int>(0), r = vector<int>(0);
    for (auto pp: p) pp & (1 << k) ? l.push_back(pp) : r.push_back(pp);
    if (l.size() > 0) bn->left = construct(l, k-1);
    if (r.size() > 0) bn->right = construct(r, k-1);
    return bn;
}

int search(int x, int k, BinaryNode* bn) {
    if (bn->data != -1) return bn->data;
    if (bn->left == nullptr) return search(x, k - 1, bn->right);
    if (bn->right == nullptr) return search(x, k - 1, bn->left);
    if (x & (1 << k)) return search(x, k - 1, bn->right);
    return search(x, k - 1, bn->left);
}

int main() {
    cin >> N >> M;
    REP(i, N) cin >> A[i];
    REP(i, M) cin >> B[i];
    
    vector<int> hoge, fuga;
    int ans = 0;

    REP2(len, 1, M+1) {
        int tmp = 0;
        REP(i, len) tmp ^= B[i];
        if (len % 2 == 0) ans = max(ans, tmp);
        if (len % 2 == 1) fuga.push_back(tmp);
        REP2(i, len, M) {
            tmp ^= B[i-len];
            tmp ^= B[i];
            if (len % 2 == 0) ans = max(ans, tmp);
            if (len % 2 == 1) fuga.push_back(tmp);
        }
    }

    REP2(len, 1, N+1) {
        int tmp = 0;
        REP(i, len) tmp ^= A[i];
        if (len % 2 == 0) ans = max(ans, tmp);
        if (len % 2 == 1) hoge.push_back(tmp);
        REP2(i, len, N) {
            tmp ^= A[i-len];
            tmp ^= A[i];
            if (len % 2 == 0) ans = max(ans, tmp);
            if (len % 2 == 1) hoge.push_back(tmp);
        }
    }


    sort(hoge.begin(), hoge.end());
    sort(fuga.begin(), fuga.end());
    hoge.erase(unique(hoge.begin(), hoge.end()), hoge.end());
    fuga.erase(unique(fuga.begin(), fuga.end()), fuga.end());
    auto at = construct(hoge, 30);
    auto bt = construct(fuga, 30);

    for (auto a: hoge) ans = max(ans, a ^ search(a, 30, bt));
    for (auto b: fuga) ans = max(ans, b ^ search(b, 30, at));
    
    cout << ans << endl;
}

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;
}