AtCoder Grand Contest 012: D - Colorful Balls

https://beta.atcoder.jp/contests/agc012/tasks/agc012_d

問題概要

N個のボールが一列に並んでおり、各ボールには色と重さが設定されている。色が同じで合計の重さがX以下のボール2つの位置を入れ替える操作と、色が異なって合計の重さがY以下のボール2つの位置を入れ替える操作という2つの操作を好きな順序で何回でも行うことができる。最終的なボールの色の並びとしてありうるものは何通りか。

N <= 2*105

解法

まず重要な点として入れ替え操作には推移的な関係がある。つまりボールAとボールBが入れ替え可能かつボールBとボールCが入れ替え可能であるならば、ボールAとボールCは(直接的には不可能でも)実質的に入れ替え可能とみなすことができる。(たとえばABCのような並びの場合、ABC->BAC->CAB->CBA のように交換していくことができ、最初と最後だけ見ればAとCを直接交換しているのと変わりがない)

このことを念頭におくと入替可能なボール同士はひとつのグループにまとめていくことができる(union-findのuniteのようなイメージ)。このようなグループ分けをすると「ひとつのグループ内でのボール移動は完全に自由」「異なるグループ間でのボール移動はありえない」ということがいえるので、あとは単なる組合せの問題になる(色a, b, cのボールがそれぞれx, y, z個ある。並び替え方は何通りか?)。

以上よりあとはグループ分けの仕方だけが問題となるが、これを効率的にやるためには各色の一番軽いボールに注目するとよい。入替の条件は重さの和なので、たとえば「色aの中で一番軽いボール」との入替が不可能ならば、当然それより重い同色のボールとの入替も不可能である。またこのように考えると、実は「異なる複数の色のボールを含みうるグループは高々1つである」ということも言える。全体の中で最も軽いボール(Aとする)の色がaだったとしたとき、a以外のボールが異なる色と繋がれるならば必ずAと繋がれることになるし、Aと同じ色のAより重いボールが異なる色と繋がれるならAも必ずそれと繋がれるからである。

というわけで答えは

  • 同色の中で最も軽いボールとの和がX以下ならグループに入れる
  • 異色の中で最も軽いボールとの和がY以下ならグループに入れる

という計算を行い、できたグループの中での色の並び替えを数えればよいということになる。

感想

最初不可能にしか見えなかったけどがんばったらできた

コード (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, std.bitmanip;

immutable long MOD = 10L^^9 + 7;

void main() {
    auto s = readln.split.map!(to!long);
    auto N = s[0].to!int;
    auto X = s[1];
    auto Y = s[2];
    auto A = new long[][](N);
    foreach (i; 0..N) {
        s = readln.split.map!(to!long);
        A[s[0].to!int-1] ~= s[1];
    }

    A = A.filter!(a => !a.empty).array;
    int M = A.length.to!int;
    
    foreach (i; 0..M) {
        A[i].sort!"a < b";
    }
    A.sort!"a[0] < b[0]";

    
    int[] B;
    
    foreach (i; 0..M) {
        if (i >= 1 && A[i][0] + A[0][0] > Y) {
            break;
        }
        
        int cnt = 1;
        long second = A.length > 1 ? A[1][0] : 1L<<59;
        
        foreach (j; 1..A[i].length) {
            if (i == 0) {
                if (A[i][0] + A[i][j] <= X || second + A[i][j] <= Y) {
                    ++cnt;
                }
            } else {
                if (A[i][0] + A[i][j] <= X || A[0][0] + A[i][j] <= Y) {
                    ++cnt;
                }
            }
        }
        B ~= cnt;
    }


    auto F = new long[](N+1);
    F[0] = 1;
    foreach (i; 0..N) {
        F[i+1] = F[i] * (i + 1) % MOD;
    }

    long ans = F[B.sum];
    foreach (b; B) {
        ans = ans * powmod(F[b], MOD-2, MOD) % MOD;
    }
    
    ans.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;
}