2023-06-06

AtCoder ABC 304 C - Virus (3Q, 灰色, 300 点)

単純な DFS / BFS 系はいよいよ灰色 diff になったのですね。

問題へのリンク

問題概要

二次元平面上に $N$ 人がいます。 $i$ 番目の人は座標 $(x_{i}, y_{i})$ にいます。

今、 $1$ 番目の人がウィルスに感染しました。ウイルスに感染した人から距離が $D$ 以内にいる人にウイルスは次々とうつります。

各人 $i$ について、ウィルスに感染したかどうかを判定してください。

制約

$1 \le N \le 2000$

考えたこと

この問題は、まず「グラフの問題だと思えるか」どうかが一つの壁になるかもしれないですね。

頂点： $N$ 人の人たち
辺：人 $i$ と人 $j$ との距離が $D$ 以下であるとき、辺 $(i, j)$ を張る

としたグラフを考えましょう。たとえば、下図のような感じになります。

こうして作ったグラフにおいて、各頂点が頂点 1 とつながっているかどうかを調べればよいです。上図の場合、頂点 2, 4, 6, 8 は Yes であり、頂点 3, 5, 7, 9 は No です。

グラフの連結性を調べる方法

グラフにおいて、頂点間がつながっているかどうかを調べる方法は、ざっくり次の 3 つの方法があります。

DFS を用いる方法
BFS を用いる方法
Union-Find を用いる方法

1, 2 の方法については、次の Qiita 記事に詳しく書いたので参考にしてもらえたらと思います。

3 の Union-Find を用いた方法については、蟻本・螺旋本・けんちょん本・鉄則本など、さまざまな本に載っています。

コード

解法 1：DFS によるコード

グラフの辺数は $O(N^{2})$ となるため、計算量は $O(N^{2})$ となります。

#include <bits/stdc++.h>
using namespace std;

// 入力データ
long long N, D;
vector<long long> X, Y;

// 頂点 i, j 間が連結かどうか
bool is_connected(int i, int j) {
    long long dis = (X[i] - X[j]) * (X[i] - X[j]) + (Y[i] - Y[j]) * (Y[i] - Y[j]);
    return dis <= D * D;
}

// DFS
void dfs(int v, vector<bool> &seen) {
    seen[v] = true;
    for (int v2 = 0; v2 < N; ++v2) {
        if (seen[v2]) continue;
        if (is_connected(v, v2)) dfs(v2, seen);
    }
}

int main() {
    // 入力
    cin >> N >> D;
    X.resize(N), Y.resize(N);
    for (int i = 0; i < N; ++i) cin >> X[i] >> Y[i];
    
    // DFS
    vector<bool> seen(N, false);
    dfs(0, seen);
    
    // 出力
    for (int v = 0; v < N; ++v) {
        cout << (seen[v] ? "Yes" : "No") << endl;
    }
}

解法 2：BFS によるコード

グラフの辺数は $O(N^{2})$ となるため、計算量は $O(N^{2})$ となります。

#include <bits/stdc++.h>
using namespace std;

// 入力データ
long long N, D;
vector<long long> X, Y;

// 頂点 i, j 間が連結かどうか
bool is_connected(int i, int j) {
    long long dis = (X[i] - X[j]) * (X[i] - X[j]) + (Y[i] - Y[j]) * (Y[i] - Y[j]);
    return dis <= D * D;
}

// 頂点 s を始点とする BFS
void bfs(int s, vector<int> &dist) {
    queue<int> que;
    dist[s] = 0;
    que.push(s);
    
    while (!que.empty()) {
        int v = que.front();
        que.pop();
        for (int v2 = 0; v2 < N; ++v2) {
            if (dist[v2] != -1 || !is_connected(v, v2)) continue;
            dist[v2] = dist[v] + 1;
            que.push(v2);
        }
    }
}

int main() {
    // 入力
    cin >> N >> D;
    X.resize(N), Y.resize(N);
    for (int i = 0; i < N; ++i) cin >> X[i] >> Y[i];
    
    // BFS
    vector<int> dist(N, -1);
    bfs(0, dist);
    
    // 出力
    for (int v = 0; v < N; ++v) {
        cout << (dist[v] != -1 ? "Yes" : "No") << endl;
    }
}

解法 3：Union-Find を用いる方法

計算量は $O(N^{2} \alpha(N))$ となります。適切な実装のもとでは $O(N^{2})$ にもなるようです (参考)。

#include <bits/stdc++.h>
using namespace std;

// Union-Find
struct UnionFind {
    // core member
    vector<int> par;

    // constructor
    UnionFind() { }
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    // core methods
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool same(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x), y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
    
    // debug
    friend ostream& operator << (ostream &s, UnionFind uf) {
        map<int, vector<int>> groups;
        for (int i = 0; i < uf.par.size(); ++i) {
            int r = uf.root(i);
            groups[r].push_back(i);
        }
        for (const auto &it : groups) {
            s << "group: ";
            for (auto v : it.second) s << v << " ";
            s << endl;
        }
        return s;
    }
};

int main() {
    // 入力
    long long N, D;
    cin >> N >> D;
    vector<long long> X(N), Y(N);
    for (int i = 0; i < N; ++i) cin >> X[i] >> Y[i];
    
    // 頂点 i, j 間が連結かどうか
    auto is_connected = [&](int i, int j) -> bool {
        long long dis = (X[i] - X[j]) * (X[i] - X[j]) + (Y[i] - Y[j]) * (Y[i] - Y[j]);
        return dis <= D * D;
    };
    
    // Union-Find
    UnionFind uf(N);
    
    // 各頂点間が連結かどうかを調べる
    for (int i = 0; i < N; ++i) {
        for (int j = i+1; j < N; ++j) {
            if (is_connected(i, j)) uf.merge(i, j);
        }
    }
    
    // 出力
    for (int i = 0; i < N; ++i) {
        cout << (uf.same(0, i) ? "Yes" : "No") << endl;
    }
}

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

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