여우가 정보섬에 올라온 이유( BOJ 17131 )

문제 : https://www.acmicpc.net/problem/17131

17131번: 여우가 정보섬에 올라온 이유

 

문제 파악하기

제시된 별들의 중 V형 별자리를 만들 수 있는 모든 경우의 수를 구하는 문제입니다. V형 별자리이라는 건 세 별 (s,t,u)가 s.x < t.x < u.x이고 s.y > t.y < u.y이면 V형 별자리라고 합니다. 결국 우리는 t보다 높은 위치에 있는 s와 u를 각각 왼쪽과 오른쪽에서 찾아야 합니다. 왼쪽에서 찾을 수 있는 s의 개수와 오른쪽에서 찾을 수 있는 u의 개수를 찾은 다음, 서로 곱해주면 t로 만들 수 있는 V형 별자리의 개수가 나옵니다. 문제는 별의 개수가 최대 200,000개라는 점이 되겠네요. 이 때, 우리는 자료구조를 사용할 수 있습니다.

 

문제 해결하기

우리가 원하는 건 현재 별(t)보다 왼쪽과 오른쪽에 각각 몇 개의 별이 V형 별자리를 만들 때 가능한지 알고 싶습니다. V형 별자리를 만들 수 있다는 의미는 t보다 y좌표 값이 항상 큰 별을 의미합니다. 이 때, 세그먼트 트리를 사용하면 효율적으로 개수를 셀 수 있습니다.

 

우선, 세그먼트 트리를 x좌표 기준으로 만듭니다. x좌표는 중복되는 값 없이 만들어야 합니다. 그리고 y좌표 순서대로 별을 하나씩 세그먼트 트리에 넣은 다음, 세그먼트 트리를 참고하여 만들 수 있는 V형 별자리의 개수를 구합니다. 현재 별보다 큰 y좌표의 개수를 세는 건 어렵지 않습니다. 왼쪽 혹은 오른쪽에 아직 들어가지 않은 별의 개수를 세면 되며, 이는 세그먼트 트리를 사용해서 지금까지 들어간 별의 개수를 카운트하여 구할 수 있습니다. 이렇듯, 정렬과 세그먼트 트리를 적절하게 사용하면 우리는 원하는 결괏값을 얻을 수 있습니다.

 

소스코드

#include <stdio.h>
#include <utility>
#include <vector>
#include <queue>
#include <algorithm>
#define NMAX 200010
#define MOD 1000000007
#define lint long long int
#define PAIR pair<int , int>
using namespace std;
 
int N, a, b;
int xsz, inpsz;
vector< PAIR > inp;
vector< int > xpos;
 
int sz;
lint segTree[1<<20];
lint cnt[NMAX];
 
lint ret;
int x, y;
queue< int > q;
 
 
// 세그먼트 트리 갱신
void update(int idx) {
    while(idx > 0) {
        segTree[idx]++;
        idx /= 2;
    }
}
 
// 세그먼트 트리 탐색
lint segSearch(int l, int r) {
    if(l>r) return 0;
    
    int lt, rt;
    lt = l+sz; rt = r+sz;
    
    lint ret=0;
    while(lt<=rt) {
        if(lt%2==1) ret += segTree[lt++];
        if(rt%2==0) ret += segTree[rt--];
        
        lt/=2; rt/=2;
    }
    
    if(l==0) return cnt[r]-ret;
    else return (cnt[r]-cnt[l-1])-ret;
}
 
// x값 위치 찾기
int bsearch(int val) {
    int l, r, mid;
    l = 0; r = xsz-1;
    
    while(l<=r) {
        mid = (l+r)/2;
        
        if(xpos[mid] == val) break;
        else if(xpos[mid] < val) l = mid+1;
        else r = mid-1;
    }
    
    return mid;
}
 
bool cmp(PAIR a, PAIR b) {
    if(a.second == b.second) return a.first<b.first;
    else return a.second<b.second;
}
 
int main() {
    //freopen("input2.txt", "r", stdin);
    
    // input
    scanf("%d", &N);
    for(int i=1;i<=N;i++) {
        scanf("%d %d", &a, &b);
        
        inp.push_back({a, b});
        xpos.push_back(a);
    }
    
    // 입력값 정렬 및 중복 제거
    sort(inp.begin(), inp.end(), cmp);
    
    // x값 추출(중복 제거)
    sort(xpos.begin(), xpos.end());
    xpos.erase(unique(xpos.begin(), xpos.end()), xpos.end());
    
    xsz = (int)xpos.size();
    inpsz = (int)inp.size();
    
    // x값 기준 누적합 계산하기
    for(int i=0;i<inpsz;i++) {
        int pos = bsearch(inp[i].first);
        cnt[pos]++;
    }
    for(int i=1;i<xsz;i++) cnt[i] += cnt[i-1];
    
    
    // Segment Tree(x값을 기준으로 만들기)
    for(sz=1;sz<xsz;sz*=2);
    for(int i=0;i<inpsz;i++) {
        x = inp[i].first; y = inp[i].second;
        int idx = bsearch(x);
        
        q.push(idx);
        update(idx+sz);
        
        // V자 별자리 경우의 수 구하기
        if(i+1==inpsz or y != inp[i+1].second) {
            while(!q.empty()) {
                int front = q.front();
                q.pop();
                
                ret = ( ret + segSearch(0, front-1)*segSearch(front+1, xsz-1) )%MOD;
            }
        }
    }
    
    printf("%lld", ret);
}
 

후기

발상 자체는 세그먼트 트리에 익숙한 사람이라면 쉽게 떠올릴 수 있습니다. 하지만 중복된 x좌표 값들이나 y좌표가 같은 별들을 한번에 저장한 다음, V형 별자리의 개수를 파악해야하는 구현이 귀찮은 문제라고 생각합니다. 세그먼트 트리의 사용법을 익히는 사람에게 추천하고 싶은 문제입니다.