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