AtCoder Regular Contest 083: F - Collecting Balls

https://beta.atcoder.jp/contests/arc083/tasks/arc083_d

問題概要

2次元平面上に2N個のロボットと2N個のボールがある。ロボットはx軸上のx座標1〜Nの箇所にN個、y軸上のy座標1〜Nの箇所にN個それぞれ設置されている。x軸上のロボットは起動されるとx=a上にあるボールでまだ残っているもののうち最もy座標が小さいものを回収し停止する。y軸上のロボットもy=b上のボールに対して同様に動作する。2N個のボールの座標が与えられたとき、ボールをすべて回収できるようなロボットの起動のさせ方の順番は何通りあるか求めよ。

N <= 105

解法

いったん順番は無視して、どのボールをどのロボットで回収するかという対応づけを考える。(x, y)にあるボールは(x, 0)のロボで取るか(0, y)のロボで取るかの2通りしかない。ここでロボットを頂点、ボールを辺としてとらえてロボ(x, 0)とロボ(0, y)をボール(x, y)が繋いでいるという形のグラフを作る。すると「ボールの回収」という事象は「ある頂点uとそこにつながっているある辺e」をひとつの組としてマッチングさせることに帰着される。ボールとロボットは1対1対応でなければならないこと、またすべてのボールを回収しなければならないことから、構築したグラフのすべての頂点と辺で過不足・重複なく組をつくることが「すべてのボールを回収する」ための必要条件である。グラフは非連結になりうるが、各連結成分どうしは順番に関して明らかに独立なので、個別に考えてよい(あとで答えを合成することも問題なくできる)。よって以降はグラフが連結であると仮定する。

ボールを取り切る、つまり頂点と隣接辺の組を過不足・重複なく作るためには、まず明らかに頂点数=辺数でなければならない。そうでないなら答えは0として終了して良い。頂点数=辺数の単純無向グラフは、木に1本辺が足された形をしている(俗に「なもりグラフ」と呼ばれることもある)。これはひとつの閉路から何本も木が生えているもの、というふうに見ることができる。ここで閉路に含まれない頂点に対応する辺は一意に定まる(葉の方から削っていくことを考えればよい)。すると閉路が残るが、閉路でのマッチングのさせかたは回る方向によって2通りありうる。これは両方考慮する必要がある。

頂点と辺の組(=ボールとロボの組)を決めたら、あとは順番を計算していく。回収においては「このボールをこのロボットに回収させるためには、事前にこのボールは回収されていなければならない」という順序の制約がある。例えば(x, 0)のロボットが(x, y)のボールを回収するためには、事前に(x, 1), (x, 2), ..., (x, y-1) にあるボールは回収されていなければいけない。今度はボールとロボの組を頂点とみなしてこのような順序関係を再びグラフにしていくと、できたグラフにおける「トポロジカル順序を守った上で並び替える方法の総数」がすなわちこの組み合わせにおけるありえる並び替えの総数となる。一般的なDAGではこのような計算は指数オーダーになってしまうが(たぶん)、この問題の条件の元でこうして作ったグラフは必ず森になるので、木DPみたいな感じで線形オーダーで総数を計算することができる。なぜ森になるかというと順序の制約が発生する場合は必ず「小さい座標のものが大きい座標のものより先」になるため。xy平面上でいうと必ず上か右に順序の矢印が向かうので、閉路は発生し得ない。

感想

解説を見て解法の概要は結構素直に飲み込めてもそれを実装に落とすところがめちゃくちゃ大変だった

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

alias Point = Tuple!(int, "x", int, "y");
alias Edge = Tuple!(int, "to", int, "x", int, "y");

immutable long MOD = 10^^9 + 7;

void main() {
    auto N = readln.chomp.to!int;
    auto G = new Edge[][](2*N); // ロボットが頂点、ボールが辺
    auto uf = new UnionFind(2*N);
    auto balls = new Point[](2*N);
    auto deg = new int[](2*N);
    auto used = new bool[](2*N);
    auto stack = new int[](2*N);
    foreach (i; 0..2*N) {
        auto s = readln.split.map!(to!int);
        int x = s[0] - 1;
        int y = s[1] - 1;
        balls[i] = Point(x, y);
        G[x] ~= Edge(y + N, x, y);
        G[y + N] ~= Edge(x, x, y);
        deg[x] += 1;
        deg[y + N] += 1;
        uf.unite(x, y + N);
    }

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

    foreach (i; 0..2*N) { // 各連結成分において 辺数 = 頂点数 が成り立ってないと駄目
        if (uf.find(i) != i) {
            continue;
        }
        if (uf.edges[i] != uf.group[i].length) {
            writeln(0);
            return;
        }
    }


    long topological_count(const long[] len, const long[] num) {
        long ret = F[len.sum];
        foreach (i; 0..len.length) {
            ret = ret * powmod(F[len[i]], MOD-2, MOD) % MOD;
            ret = ret * num[i] % MOD;
        }
        return ret;
    }

    long calc(Tuple!(int, Edge)[] matching) {
        int n = matching.length.to!int;

        Tuple!(int, int)[int] X;
        Tuple!(int, int)[int] Y;

        auto tree = new int[][](n);
        auto is_root = new bool[](n);
        fill(is_root, true);

        foreach (i; 0..n) {
            if (matching[i][0] < N) {
                X[matching[i][1].x] = tuple(matching[i][1].y, i);
            } else {
                Y[matching[i][1].y] = tuple(matching[i][1].x, i);
            }
        }

        // ひとつのマッチングをひとつの頂点とみなして、マッチングの順序関係からDAGをつくる(必ず有向森になる)

        foreach (i; 0..n) {
            if (matching[i][0] < N) {
                if (matching[i][1].y in Y && matching[i][1].x < Y[matching[i][1].y][0]) {
                    tree[Y[matching[i][1].y][1]] ~= i;
                    is_root[i] = false;
                }
            } else {
                if (matching[i][1].x in X && matching[i][1].y < X[matching[i][1].x][0]) {
                    tree[X[matching[i][1].x][1]] ~= i;
                    is_root[i] = false;
                }
            }
        }

        // つくった森をトポロジカルソートするやり方が何通りあるか数える

        Tuple!(long, long) dfs(int n, int p) {
            long[] len, num;
            foreach (m; tree[n]) {
                if (m == p) continue;
                auto t = dfs(m, n);
                len ~= t[0];
                num ~= t[1];
            }
            return tuple(len.sum+1, topological_count(len, num));
        }

        long[] len, num;
        foreach (i; 0..n) {
            if (!is_root[i]) continue;
            auto t = dfs(i, -1);
            len ~= t[0];
            num ~= t[1];
        }

        return topological_count(len, num);
    }

    long solve(const int[] nodes) {
        int[] cycle_nodes;
        Edge[] cycle_edges;
        Tuple!(int, Edge)[] not_cycle;

        int sp = -1;
        foreach (n; nodes) if (deg[n] == 1) stack[++sp] = n;

        while (sp >= 0) {
            int n = stack[sp--];
            used[n] = true;
            foreach (e; G[n]) {
                if (used[e.to]) continue;
                deg[e.to] -= 1;
                if (deg[e.to] == 1) stack[++sp] = e.to;
                not_cycle ~= tuple(n, e);
                break;
            }
        }

        foreach (start; nodes) {
            if (used[start]) continue;
            int next = start;
            while (cycle_nodes.empty || next != start) {
                cycle_nodes ~= next;
                int cur = next;
                used[cur] = true;
                foreach (e; G[cur]) {
                    if (e.to == start && cycle_nodes.length == 2) continue;
                    if (e.to != start && used[e.to]) continue;
                    cycle_edges ~= e;
                    next = e.to;
                    break;
                }
            }
            break;
        }


        Tuple!(int, Edge)[] cycle1; // サイクル順周り
        Tuple!(int, Edge)[] cycle2; // サイクル逆周り
        int cnt = cycle_nodes.length.to!int;

        foreach (i; 0..cnt) {
            cycle1 ~= tuple(cycle_nodes[i], cycle_edges[i]);
        }
        foreach (i; 0..cnt) {
            cycle2 ~= tuple(cycle_nodes[(i+1)%cnt], Edge(cycle_nodes[i], cycle_edges[i].x, cycle_edges[i].y));
        }

        return (calc(cycle1 ~ not_cycle) + calc(cycle2 ~ not_cycle)) % MOD;
    }


    long[] len, num;

    foreach (i; 0..2*N) {
        if (uf.find(i) != i) {
            continue;
        }
        len ~= uf.group[i].length.to!long;
        num ~= solve(uf.group[i]);
    }

    long ans = topological_count(len, num);
    ans.writeln;
}


class UnionFind {
    int N;
    int[] table;
    int[][] group;
    int[] edges;

    this(int n) {
        N = n;
        table = new int[](N);
        fill(table, -1);
        group = N.iota.map!(i => [i]).array;
        edges = new int[](N);
    }

    int find(int x) {
        return table[x] < 0 ? x : (table[x] = find(table[x]));
    }

    void unite(int x, int y) {
        x = find(x);
        y = find(y);
        if (x == y) {
            edges[x] += 1;
            return;
        }
        if (table[x] > table[y]) swap(x, y);
        table[x] += table[y];
        table[y] = x;
        group[x] ~= group[y].dup;
        group[y] = [];
        edges[x] += edges[y] + 1;
    }
}

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