[Offline Dynamic Connectivity] BOJ 15316 현수 시티

https://www.acmicpc.net/problem/15316

15316번: 현수시티

문제 요약 : 정점과 간선이 주어지고 쿼리가 주어지는데 i번째 간선을 연결했다가 끊었다가할 때 어떤 정점 u, v가 연결되어있는지 확인하는 문제이다.

먼저 u,v 가 연결되어있는지를 확인하는 작업은 union find 를 사용한다.

그런데 끊어야 하기 때문에 path compression 은 사용할 수 없고 rank compression만 사용해야한다.

그리고 특정 작업을 롤백하기 위해서 합쳐질 때의 정점과 정점의 rank 변화량을 스택에 저장해주면 된다.

#define N 200001
int p[N];
int r[N];
struct trip {
    int x, y, c;
};
stack<trip>st;

parent 와 rank, 그리고 스택을 선언했다.

find 함수에서path compression을 쓰면

int find(int x) {
    if (x == p[x]) return x;
    return p[x] = find(p[x]);
}

로 구현하지만 이렇게 되면 롤백이 불가능하므로

int find(int x) {
    if (x == p[x]) return x;
    return find(p[x]);
}

로 구현한다.

int merge(int x, int y) {
    x = find(x), y = find(y);
    if (x == y) return 0;
    if (r[x] < r[y]) swap(x, y);
    p[y] = x;
    if (r[x] == r[y]) {
        r[x]++;
        st.push({ x,y,1 });
    }
    else st.push({ x,y,0 });
    return 1;
}
void rollback(int cnt) {
    while (cnt--) {
        trip e = st.top(); st.pop();
        p[e.y] = e.y;
        if (e.c) r[e.x]--;
    }
}

x,y 를 merge할때 rank가 변했음을 같이 저장해주고 있다.
rollback은 스택에 있는 것을 꺼내서
원래대로 되돌려주는 작업만을 하면 된다.

 

그 다음 중요한 것은 어떤 데이터들을 어떻게 저장할 것인가가 필요하다.
우선 edge들이 들어오므로 edge를 배열로 저장해야 할 것이다.
그리고 정점이 연결되어있는지에 대한 쿼리를 처리해야하는데 이 때 쿼리를 위한 정점도 저장해두자.
그리고 어떤 간선이 사용되고 있는지, 그리고 그 간선on/off 쿼리를 위해 간선의 상태가 변할 때를 저장해줄 배열이 하나 필요하다.

pair<int, int>ans[N]; // 2번 쿼리에 대한 정점을 저장
pair<int, int>edge[N]; //edge 를 저장
int queryNum[N]; // Q개의 쿼리가 들어오는데 2번쿼리가 들어올 때마다 1씩 값을 증가시켜서 어떤 1번쿼리가 몇 번째의 2번 쿼리에 영향을 줄지 계산하기 위한 배열
int light[N]; //어떤 간선이 on / off 되었는지를 나타내는 배열
int before[N]; // 어떤 간선이 on 되는 시점을 저장하는 배열

위와 같이 선언해준다.
그 다음은 간선들이 i번째 2번쿼리부터 on이고 j번째 2번쿼리에 off되었을 때의 수명을 저장하기 위한 배열이 필요하다.

struct Life {
    int idx, s, e; // idx번째 간선은 2번째 쿼리 s ~ e 까지의 구간에 사용될 것이다.
};
vector<Life> lives;

1번 쿼리가 몇 번들어올지 모르므로 동적으로 추가할 수 있는 vector로 선언했다.

마지막으로 세그먼트트리가 필요하다.
2번쿼리만을 생각해보면 2번쿼리가 x개 들어온다면

2번쿼리 넘버링을 1 2 ... x 로 할 수 있을 것이고 이것을 구간으로 생각한다면
1번 쿼리가 2번의 특정 구간에 영향을 주는 것으로 생각할 수 있다.

 

이제 우리는 x를 구간으로 삼는 세그먼트트리를 만들면 된다.
세그먼트트리의 한 노드는 특정 구간을 담당하게 만들어줄 수 있다.
즉 이 구간이 위에 선언한 Life 의 s, e에 잘 매칭해주면 된다.
우리는 특정 구간(i번째 2번쿼리 ~ j번째 2번쿼리)을 담당하는 세그먼트트리의 노드에 해당할 간선들을 모두 저장해주면 해결할 수 있다.

그렇기 때문에 세그먼트 트리는

vector<int> seg[4 * N];

로 선언하고
update 와 query 함수를 작성해주자.

void update(int node, int s, int e, int i, int j, int val) {
    if (e<i || j < s) return;
    if (i <= s && e <= j) {
        seg[node].push_back(val);
        return;
    }
    int m = s + e >> 1;
    update(node * 2, s, m, i, j, val);
    update(node * 2 + 1, m + 1, e, i, j, val);
}

void query(int node, int s, int e) {
    int count = 0;
    for (int i : seg[node]) count += merge(edge[i].first, edge[i].second);
    if (s == e) {
        if (find(ans[s].first) == find(ans[s].second)) cout << "YES\n";    
        else cout << "NO\n";
        rollback(count);
        return;
    }
    int m = s + e >> 1;
    query(node * 2, s, m);
    query(node * 2 + 1, m + 1, e);
    rollback(count);
}

update에서는 구간에 해당하는 간선 번호(val)를 모두 저장하게 해두었고
query에서는 노드의 하위구간에 영향을 미칠 친구들을 모두 merge한 후 rollback시켜주고 있다.

이제 메인을 작성해보자

int main() {
    ios_base::sync_with_stdio(false), cin.tie(0);
    int n, m; cin >> n >> m;
    for (int i = 1; i <= n; ++i) p[i] = i;
    for (int i = 1, u, v; i <= m; ++i) {
        cin >> edge[i].first >> edge[i].second;
    }
    for (int i = 1; i <= m; ++i) {
        cin >> light[i];
    }
    int Q; cin >> Q;
    for (int i = 1, cmd, u, v; i <= Q; ++i) {
        cin >> cmd;
        if (cmd == 1) {
            cin >> u;
        }
        else {
            cin >> u >> v;
        }
    }
}

유니온파인드를 위한 p를 초기화하고 입력을 받는 부분이다.
cmd == 1 이면 쿼리카운트는 증가시킬 필요가 없고
cmd==2라면 쿼리카운트는 증가시켜야 한다.
이를 추가해주자.

if (cmd == 1) {
    cin >> u;
    queryNum[i] = queryNum[i - 1];
}
else {
    cin >> u >> v;
    queryNum[i] = queryNum[i - 1] + 1;
}

그 다음은 cmd==1일 때 light[u] 가 on이라면 off시켜주고 off 라면 on시켜주면서 on된 시점을 저장해주면 된다.
그리고 cmd==2라면 해답을 구하기위한 정점들을 저장해주자.

    for (int i = 1, cmd, u, v; i <= Q; ++i) {
        cin >> cmd;
        if (cmd == 1) {
            cin >> u;
            queryNum[i] = queryNum[i - 1];
            if (light[u]) {
                light[u] = 0;
            }
            else {
                light[u] = 1;
                before[u] = i;
            }
        }
        else {
            cin >> u >> v;
            queryNum[i] = queryNum[i - 1] + 1;
            ans[queryNum[i]] = { u,v };
        }
    }

이제 queryNum이 저장된 형태를 살펴보면
1 1 1 2 2 2 2 3 3 3 3 과 같은 형태일 것이다.
굵게 표시된 부분은 2번쿼리가 실행되었을 때이므로 굵게 표시된 이후의 기본 숫자들은 모두 1번쿼리가 실행된 것이다.
이를 통해 만약 1번쿼리가 들어온 시점x의 queryNum[x] 가 2라고 생각해보면 이미 2번쿼리는 실행되고 난 이후일 것이다. 우리는 before 배열에 on되는 시점을 저장해두었으므로 queryNum[before[y]] 는 이미 실행된 쿼리이므로 1을 더해주면 된다. 즉 queryNum[before[y]]이 2였다면 3번 쿼리부터 영향을 주게 될 것이라는 뜻이다.
그럼 light[u]가 on일 때 1번쿼리가 들어왔다면
이제 끄는 시점이 된 것이고
시작지점은 queryNum[before[u]] + 1이 되는 것이고
끝 지점은 queryNum[i] 가 되는 것이다.
이 것이 Life를 이루게 될 것이다.

            if (light[u]) {
                light[u] = 0;
                int start = queryNum[before[u]] + 1;
                int end = queryNum[i];
                lives.push_back({ u, start, end });
            }

와 같이 될 것이다.
즉 현재 상태를 보면

    for (int i = 1, cmd, u, v; i <= Q; ++i) {
        cin >> cmd;
        if (cmd == 1) {
            cin >> u;
            queryNum[i] = queryNum[i - 1];
            if (light[u]) {
                light[u] = 0;
                int start = queryNum[before[u]] + 1;
                int end = queryNum[i];
                lives.push_back({ u, start, end });
            }
            else {
                light[u] = 1;
                before[u] = i;
            }
        }
        else {
            cin >> u >> v;
            queryNum[i] = queryNum[i - 1] + 1;
            ans[queryNum[i]] = { u,v };
        }
    }

일 것이다.
그런데 잘 생각해보면 끝까지 켜져있는 경우가 있을 수 있다. 이것은 구간 끝까지 살아서 영향을 주는 간선이 될 것이므로 이 것을 처리해 주자.

for (int i = 1; i <= m; ++i) if (light[i]) lives.push_back({ i, queryNum[before[i]] + 1, queryNum[Q] });

그럼 마지막으로 모아둔 모든 life를 이용해서 세그먼트트리를 업데이트 해주자.

for (Life& l : lives) update(1, 1, queryNum[Q], l.s, l.e, l.idx);

세그먼트트리노드는 1부터 시작하며 1~queryNum[Q] 의 구간을 다루고 있으며 해당 라이프의 구간 l.s, l.e와 간선의 번호 l.idx를 업데이트 해주면 된다.
마지막엔 query를 한번만 실행해주면 된다.

query(1, 1, queryNum[Q]);

이유는 세그먼트트리는 이진트리로 모든 노드를 탐색하면되고 이 때 순서가 왼쪽이 번호가 작은 구간이 들어있기 때문에 리프노드에 도달했을 때 2번쿼리의 답을 차례대로 출력해줄 것이다.

코드최종

#include <iostream>
#include <vector>
#include <map>
#include <stack>
using namespace std;
#define N 200001
int p[N];
int r[N];
struct trip {
    int x, y, c;
};
stack<trip>st;
int find(int x) {
    if (x == p[x]) return x;
    return find(p[x]);
}
int merge(int x, int y) {
    x = find(x), y = find(y);
    if (x == y) return 0;
    if (r[x] < r[y]) swap(x, y);
    p[y] = x;
    if (r[x] == r[y]) {
        r[x]++;
        st.push({ x,y,1 });
    }
    else st.push({ x,y,0 });
    return 1;
}
void rollback(int cnt) {
    while (cnt--) {
        trip e = st.top(); st.pop();
        p[e.y] = e.y;
        if (e.c) r[e.x]--;
    }
}
pair<int, int>ans[N];
pair<int, int>edge[N];
int queryNum[N];
int light[N];
int before[N];
struct Life {
    int idx, s, e;
};
vector<Life> lives;
vector<int> seg[4 * N];
void update(int node, int s, int e, int i, int j, int val) {
    if (e<i || j < s) return;
    if (i <= s && e <= j) {
        seg[node].push_back(val);
        return;
    }
    int m = s + e >> 1;
    update(node * 2, s, m, i, j, val);
    update(node * 2 + 1, m + 1, e, i, j, val);
}

void query(int node, int s, int e) {
    int count = 0;
    for (int i : seg[node]) count += merge(edge[i].first, edge[i].second);
    if (s == e) {
        if (find(ans[s].first) == find(ans[s].second)) cout << "YES\n";    
        else cout << "NO\n";
        rollback(count);
        return;
    }
    int m = s + e >> 1;
    query(node * 2, s, m);
    query(node * 2 + 1, m + 1, e);
    rollback(count);
}
int main() {
    ios_base::sync_with_stdio(false), cin.tie(0);
    int n, m; cin >> n >> m;
    for (int i = 1; i <= n; ++i) p[i] = i;
    for (int i = 1; i <= m; ++i) {
        cin >> edge[i].first >> edge[i].second;
    }
    for (int i = 1; i <= m; ++i) {
        cin >> light[i];
    }
    int Q; cin >> Q;
     for (int i = 1, cmd, u, v; i <= Q; ++i) {
        cin >> cmd;
        if (cmd == 1) {
            cin >> u;
            queryNum[i] = queryNum[i - 1];
            if (light[u]) {
                light[u] = 0;
                int start = queryNum[before[u]] + 1;
                int end = queryNum[i];
                lives.push_back({ u, start, end });
            }
            else {
                light[u] = 1;
                before[u] = i;
            }
        }
        else {
            cin >> u >> v;
            queryNum[i] = queryNum[i - 1] + 1;
            ans[queryNum[i]] = { u,v };
        }
    }
    for (int i = 1; i <= m; ++i) if (light[i]) lives.push_back({ i, queryNum[before[i]] + 1, queryNum[Q] });
    for (Life& l : lives) update(1, 1, queryNum[Q], l.s, l.e, l.idx);
    query(1, 1, queryNum[Q]);
}