けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 076 D - Built? (500 点)

ゆかたゆたんと一緒に解いたん。
今後まだ解いてない様々な 500 点問題について何がポイントになっているのかをブログ書きながら明らかにして行きたい。500 点問題の苦手意識を克服するぞい。

問題へのリンク

問題概要

二次元平面上に  N 個の点が与えられる。これらの各点を頂点にもつ完全グラフを考える。座標  (a, b) な頂点と座標  (c, d) な頂点とを結ぶ辺のコストは  \min(|a-c|, |b-d|) で与えられる。

このグラフの最小全域木の重みを求めよ。

制約

  •  2 \le N \le 10^{5}

考えたこと

最小全域木 (MST) を求める問題!!!

しかしそのまんま Kruskal 法とか Prim 法を適用したいなと思うと、頂点数が  O(N) で、枝数が  O(N^{2}) あるため、 O(N^{2}\log{N}) の計算量を要してしまう!
(そもそも、 O(N^{2}) の枝をもつとメモリに収まらない)

こうした密グラフな問題は MST に限らず、考えるべき辺数を減らせないか考えるのはすごくあるあるな流れな気がする。特に MST だとすごくあるあるな感じ。

今回、x 座標や y 座標とかで遠く離れた点同士はあまり結ばれなさそうなイメージがある。実際

  • x 座標順にソート
  • y 座標順にソート

という 2 通りの並びを考えたとき、どちらの基準でも隣接していない辺は考える必要がないことが示せる。

キーエンス プログラミング コンテスト 2019 E - Connecting Cities と同様の考察ができる。

u と v の距離は x 座標方向と y 座標方向のうち x 座標の方が短いとする。ここで、u と v との間に頂点 w が挟まっていると

uv の重み 
= uv の x 座標差
> uw の x 座標差
>= uw の重み

となる。同様に「uv の重み > vw の重み」も示せる。よって、uv の重みは、uw や vw の重み以上なので、辺 uv は考える必要がなくなる。

そのようにして考えるべき辺の候補を絞ると辺数が  O(N) になるので、Kruskal 法などで  O(N\log{N}) になって間に合う。

#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
using namespace std;

struct UnionFind {
    vector<int> par;
    UnionFind(int n) : par(n, -1) { }
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    void merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return;
        if (-par[x] < -par[y]) swap(x, y);
        par[x] += par[y];
        par[y] = x;
    }
};
    
int main() {
    int N; cin >> N;
    vector<long long> x(N), y(N);
    for (int i = 0; i < N; ++i) cin >> x[i] >> y[i];

    // Kruskal 法で使用する枝
    using pint = pair<int,int>;
    using Edge = pair<long long, pint>;
    vector<Edge> edges;

    // x 軸方向にソートして隣接している部分のみ
    vector<int> ids(N);
    iota(ids.begin(), ids.end(), 0);
    sort(ids.begin(), ids.end(), [&](int i, int j) { return x[i] < x[j]; });
    for (int i = 0; i+1 < ids.size(); ++i) {
        int u = ids[i], v = ids[i+1];
        edges.push_back(Edge(x[v] - x[u], pint(u, v)));
    }

    // y 軸方向にソートして隣接している部分のみ
    sort(ids.begin(), ids.end(), [&](int i, int j) { return y[i] < y[j]; });
    for (int i = 0; i+1 < ids.size(); ++i) {
        int u = ids[i], v = ids[i+1];
        edges.push_back(Edge(y[v] - y[u], pint(u, v)));
    }

    // Kruskal 法
    sort(edges.begin(), edges.end());
    UnionFind uf(N);
    long long res = 0;
    for (auto e : edges) {
        int u = e.second.first, v = e.second.second;
        long long cost = e.first;
        if (uf.issame(u, v)) continue;
        uf.merge(u, v);
        res += cost;
    }
    cout << res << endl;
}