[Codeforces] Number of Groups(UnionFind, Greedy)

https://codeforces.com/contest/1691/problem/E

Problem - E - Codeforces

이번 문제 풀이의 가장 큰 풀이는 크게 두 가지로 나눈다.

 

1. 그룹을 만든다.

2. 그룹의 수를 센다.

 

처음에는 크게 1번의 경우에는 Greedy 알고리즘, 2번의 경우에는 Union-Find 알고리즘이 쓰이면 좋을 것 같다는 떠오른다.

 

starting and ending points를 set로 저장한다. 

예를 들어 (0, 10)와 (11,12) 가 있다면 우리는 0, 10, 11, 12를 색상 상관없이 set에 저장한다. 그리고 이 점들을 오름차순으로 정렬한다.

 

우리는 색상별로 2개의 set를 가진다.

이 sets에서 starting index에는 도달했으나 ending index에는 도달하지 못한 segments를 저장한다.

즉 시작했으나 아직 끝나지 않은 세그먼트를 뜻한다.

 

만약 point $x$에 있다면

- 만약 세그먼트의 시작점에 해당된다면

1. 색상에 해당하는 set에 해당 세그먼트의 ending point를 추가한다.

2. 다른 색상에 해당하는 set과 매치되는 세그먼트를 모두 DSU 머지(Disjoint set union merge)한다.

3. 다른 색상의 set에서 가장 큰 ending point를 가진 세그먼트 하나를 제외한 모든 세그먼트를 삭제한다.

 

- 만약 세그먼트의 끝점에 해당된다면 

해당 컬러의 set에서 세그먼트를 삭제

 

왜 가장 큰 ending point를 가진 세그먼트를 제외하고 모든 세그먼트는 제거해도 될까?

조금 더 생각해보면 간단한 문제였다.

이 문제는 그룹 기반이기 때문에 가장 큰 ending point를 가진 세그먼트도 해당 그룹에 연결되어 있기 때문에 뒤에 있는 세그먼트들은 다른 세그먼트와 연결되지 않았더라도 가장 큰 ending point를 가진 세그먼트만 연결되면 해당 그룹으로 속해지기 때문이다.

따라서 그리디하게 문제를 풀 수 있게 된다.

 

그리고 linear하게 탐색을 할 때 일일이 점을 방문하는 것이 아니라 시작점, 끝점만 저장해서 상태가 달라지는 경우만 방문하도록 하였다.

 

#include <vector>
#include <iostream>
#include <tuple>
#include <algorithm>
#include <set>
typedef long long ll;
using namespace std;
struct Point {
    int color, p1, p2, id;
    bool closed;
};

// p1 오름차순 정렬
bool cmp(const Point a, const Point b) {
    if(a.p1 == b.p1) return a.closed < b.closed;
    return a.p1 < b.p1;
}
// Union-Find
int parent(vector<int> &parents, int i) {
    if(parents[i] == i) return i;
    return parents[i] = parent(parents, parents[i]);
}

void unify(vector<int> &parents, vector<int> &sizes, int a, int b) {
    int pa = parent(parents, a);
    int pb = parent(parents, b);
    
    if(sizes[a] < sizes[b]) swap(pa, pb);
    if(pa != pb) {
        parents[pb] = pa;
        sizes[pa] += sizes[pb];
    }
}

// group 개수 세기
int cnt(vector<int> &parents) {
    int ct = 0;
    for(int i = 0;i < parents.size();i++) {
        if(parents[i] == i) {
            ct++;
        }
    }
    return ct;
}

int main() {
   int t;
   cin >> t;
   while (t--) {
       int n, c, l, r;
       vector<Point> points;
       cin >> n;
       vector<int> parents(n);
       vector<int> sizes(n);
       for(int i = 0;i < n;i++) {
           cin >> c >> l >> r;
           
           points.push_back({c, l, r, i, false});
           points.push_back({c, r, l, i, true});
       }
       
       sort(points.begin(), points.end(), cmp);
       
       for(int i = 0;i < n;i++) {
           parents[i] = i;
       }
       // {r, id}
       // set이므로 자동으로 r기준 오름차순되어 삽입
       vector<set<pair<int, int>>> open(2);
       
       for(auto &p : points){
           // 해당 컬러의 set에서 세그먼트 제거
           if(p.closed) open[p.color].erase({p.p1, p.id});
           else {
               open[p.color].insert({p.p2, p.id});
               // 다른 색상의 set에 남아있는세그먼트와 합치기(아직 종료된 것이 아님이 확실)
               // 합친 후 해당 set에서 제거(ending points가 가장 먼 세그먼트 제외)
               // 어차피 ending points가 가장 먼 세그먼트와도 합치면 다른 세그먼트와도 연결되는 것이기 때문
               while(open[p.color ^ 1].size() > 1) {
                   auto p2 = *open[p.color ^ 1].begin();
                   int r = p2.first;
                   int id = p2.second;
                   unify(parents, sizes, p.id, id);
                   open[p.color ^ 1].erase({r, id});
               }
               // 다른 색상의 set에  세그먼트가 하나라면 현재 세그먼트와 합치기만 하고 제거는 안 함
               if((open[p.color ^ 1].size() == 1)) {
                   unify(parents, sizes, p.id, open[p.color ^ 1].begin() -> second);
               }
           }
       }
    
       cout << cnt(parents) << '\n';
   }
}