srupのメモ帳

競プロで解いた問題や勉強したことを記録していくメモ帳

yukicoder No.416 旅行会社

問題

問題概要

省略

解法

union findは辺を削除することができない。問題は、クエリを与えられた順にみると、辺を削除するような形だが、クエリを逆からみれば、辺をつなげていく形になる。このようになると、union findで管理できる。
また今回学んだunion findはデータ構造をマージする一般的テクを用いたもの。これは、リストでグループを管理しながら、uniteするときは、小さいグループを大きいものへつなげるという工夫で、uniteするときの計算量がならし O(log n) になるというもの。リストで管理しているので、同じグループに属しているものを知りたいときに便利??
この問題をrootで管理するUFで解こうと思ったけど、同じグループに属しているものを調べるときに、すべての頂点を調べる方法しかできなかった。これではTLEしてしまう。

//TLEしてしまう部分
for (int j = 1; j < n; ++j){
    if(ans[j] == 0 && ds.same(0, j)){
        ans[j] = i + 1;
    }
}

工夫すればできるはず。

topcoder.g.hatena.ne.jp

ミス

後ろから見ていくタイプのunion findだということはするにわかったが、rootで管理するunion findを使い、新たに、辺をつなげた時に、新しく頂点0と同じグループに属しているものを調べるときに、頂点分だけ調べる方法しか思いつかなかった。その部分でTLE。

コード

データ構造をマージする一般的テクを利用したUFは、同じグループに属する頂点をリストで管理しているため、同じグループに属しているものを調べるときに、n個の頂点を調べなくていい。
AC

#include <iostream>
#include <algorithm>
#include <vector>
#include <set>
using namespace std;
typedef long long ll;
#define rep(i,n) for(int i=0;i<(n);i++)

//Quick Find Weighted
const int MAX_N = 200010;
int i2g[MAX_N];//i2g[i] := 頂点iが所属するグループ
vector<int> g2i[MAX_N];//g2i[g]:= グループgに所属する頂点番号

// アイテムiaとアイテムibは同じグループに所属しているか?
bool issame(int ia, int ib){
    return i2g[ia] == i2g[ib];
}

// アイテムiaの所属するグループとアイテムibの所属するグループを1つにする(異なるグループに属するものに)
void merge(int ia, int ib){
    if(issame(ia, ib)) return;
    //iaの所属するグループがibの所属するグループより小さくならないようにする(一般的マージテク)
    if(g2i[i2g[ia]].size() < g2i[i2g[ib]].size()){
        swap(ia, ib);
    }
    int ga = i2g[ia], gb = i2g[ib];//グループgbの方が要素数が少ない
    for(auto u : g2i[gb]){//uには頂点の番号はいる
        i2g[u] = ga;//グループの番号を更新
    }
    g2i[ga].insert(g2i[ga].end(), g2i[gb].begin(), g2i[gb].end());//つなげる
    g2i[gb].clear();
}

void init(int n){
    for (int i = 0; i < n; ++i){
        i2g[i] = i;
        g2i[i].push_back(i);
    }
}

int ans[MAX_N];

int main(void){
    int n, m, q; cin >> n >> m >> q;
    vector<pair<int, int> > cd;
    set<pair<int, int> > s;//最後まで残っている橋を入れる

    for (int i = 0; i < m; ++i){
        int a, b; cin >> a >> b;
        a--; b--;
        s.insert(make_pair(a, b));
    }
    for (int i = 0; i < q; ++i){
        int c, d; cin >> c >> d;
        c--; d--;
        cd.push_back(make_pair(c, d));
        s.erase(make_pair(c, d));
    }

    init(n);//初期化
    for(auto u : s){//最後まで残っている橋をつなげる
        merge(u.first, u.second);
    }

    rep(i, n) ans[i] = 0;
    int g = i2g[0];//最後まで0と同じグループにあるものは、-1
    for(int v : g2i[g]){
        ans[v] = -1;
    }

    for (int i = q - 1; i >= 0 ; --i){
        int a = cd[i].first, b = cd[i].second;
        if(!issame(a, b)){
            if(!issame(0, a) && !issame(0, b)){//0につながることはない
                merge(a, b);
                continue;
            }

            if(issame(0, a)){//0とaが同じグループなら、bも0と同じグループになる
                for(auto u : g2i[i2g[b]]){
                    if(ans[u] == 0) ans[u] = i + 1;
                }
            }else{//0とbが同じグループなら、aも0と同じグループになる
                for(auto u : g2i[i2g[a]]){
                    if(ans[u] == 0) ans[u] = i + 1;
                }
            }
            merge(a, b);
        }
    }
    for (int i = 1; i < n; ++i){
        printf("%d\n", ans[i]);
    }
    return 0;
}

rootで管理するUFをそのまま使うと、グループに所属するものを列挙するときに、時間がかかり、TLE。 工夫すれば行けるはず。
TLE

#include <iostream>
#include <algorithm>
#include <vector>
#include <set>
using namespace std;
typedef long long ll;
#define rep(i,n) for(int i=0;i<(n);i++)

class DisjointSet{
    public:
        vector<int> rank, p;//rank:木の高さ p:親の頂点番号
        DisjointSet(){}
        DisjointSet(int size){//頂点の数
            rank.resize(size, 0);
            p.resize(size, 0);
            rep(i, size) makeSet(i);
        }
        void makeSet(int x){
            p[x] = x;
            rank[x] = 0;
        }
        //同じ木にあるか判定(今回使わない)
        bool same(int x, int y){
            return findSet(x) == findSet(y);
        }
        // 木どうしをくっつける
        void unite(int x, int y){
            link(findSet(x), findSet(y));
        }
        //木の高さを考慮して木どうしをくっつける
        void link(int x, int y){
            if(rank[x] > rank[y]){
                p[y] = x;
            }else{
                p[x] = y;
                if(rank[x] == rank[y]) rank[y]++;
            }
        }
        //親を探す(ルートまで)
        int findSet(int x){
            if(x != p[x]){
                p[x] = findSet(p[x]);
            }
            return p[x];
        }
};

int main(void){
    int n, m, q; cin >> n >> m >> q;
    vector<pair<int, int> > cd;
    set<pair<int, int> > s;//最後まで残っている橋を入れる

    for (int i = 0; i < m; ++i){
        int a, b; cin >> a >> b;
        a--; b--;
        s.emplace(make_pair(a, b));
    }
    for (int i = 0; i < q; ++i){
        int c, d; cin >> c >> d;
        c--; d--;
        cd.push_back(make_pair(c, d));
        s.erase(make_pair(c, d));
    }

    DisjointSet ds = DisjointSet(n);
    for(auto u : s){
        ds.unite(u.first, u.second);
    }

    vector<int> ans(n, 0);
    for (int i = 1; i < n; ++i){
        if(ds.same(0, i)){
            ans[i] = -1;
        }
    }

    for (int i = q - 1; i >= 0; --i){
        int a = cd[i].first, b = cd[i].second;
        if(!ds.same(a, b)){
            if(!ds.same(0, a) && !ds.same(0, b)){//0につながることはない
                ds.unite(a, b);
                continue;
            }

            ds.unite(cd[i].first, cd[i].second);
            //根を保持して行うunion findは1グループが木構造で管理されていて、リストで管理されないため、
            //グループに所属するものを列挙するのに、毎回頂点数分なめないといけない??
            //どの頂点が 0 と同じグループになったかを確かめるのに、毎回すべての頂点に対して、sameを行わないといけない?
            for (int j = 1; j < n; ++j){
                if(ans[j] == 0 && ds.same(0, j)){
                    ans[j] = i + 1;
                }
            }
        }
    }
    for (int i = 1; i < n; ++i){
        printf("%d\n", ans[i]);
    }
    return 0;
}