백설공주와 난쟁이( BOJ 2912 )

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

2912번: 백설공주와 난쟁이

 

문제 파악하기

구간 [a, b]사이의 숫자 중 절반 이상의 개수를 가진 숫자가 있는지 없는지 판단하는 문제입니다. 데이터의 개수는 총 300,000개이며 쿼리의 개수는 10,000개이기에 각 쿼리 당 효율적인 알고리즘이 요구되는 문제입니다. 이 문제를 풀기 위해서는 우선, 절반 이상의 개수를 가지기 위해서 필요한 조건을 생각해보아야 합니다.

 

문제 해결하기

우리가 N개의 데이터를 동일한 크기로 묶음을 지어 나눈 다음, 각각의 묶음에서 개수가 가장 많은 숫자를 대표 숫자라고 하겠습니다. 그럼 구간 [a, b]에서 만약 절반 이상의 개수를 가진 숫자가 있다면 그 숫자는 구간 [a, b] 안에 포함된 대표 숫자 중 하나입니다. 이는 당연한 이야기인게 절반 이상의 개수를 가지기 위해서는 최소 1개 이상의 묶음에서 절반 이상의 개수를 가지고 있어야 합니다. 따라서, 정답을 구하기 위해서는 구간 [a, b]에 포함된 대표 숫자들을 뽑은 다음, 구간 [a, b]에 총 몇 개가 있는지 세어보면 됩니다. 만약 어떤 대표 숫자도 절반 이상의 값을 가지지 못한다면 no를 출력하고, 어떤 대표 숫자가 절반 이상의 값을 가진다면 yes를 출력하면 됩니다.

 

예를 들어 16개의 숫자를 다음과 같이 묶어봅시다. 그리고 각각의 숫자에 대한 대표 숫자를 표현하면 다음과 같습니다.

이렇게 만든 다음, 구간 [1, 16]에서 절반 이상의 개수를 가지고 있을 수 있는 숫자의 후보는 {1, 2} 총 2개입니다. 우리는 이렇게 후보를 뽑은 다음에 구간[1, 16]에서 1의 개수와 2의 개수를 세서 절반 이상의 개수를 가지고 있는지 판단하면 됩니다.

 

그럼 N개의 데이터를 적당한 묶음으로 나눠봅시다. 여기서 묶을 수 있는 방법은 여러 가지 있습니다. 머지소트트리의 형태로 묶을 수도 있으며, 위의 그림과 같이 √N개씩 묶을 수도 있습니다. 머지소트트리로 묶는다면 후보 대표 숫자를 탐색하는데 O(logN)이 소요되며, √N개씩 묶는다면 탐색 시 O(√N)이 소요됩니다. 다행히 300,000개의 데이터를 탐색하는데에는 어느 구조든 상관없습니다.

 

이렇게 적절하게 데이터를 묶은 후, 대표 숫자를 구한 다음에는 구간[a, b]에서 각각의 대표 숫자 개수를 파악해야 합니다. 우리는 여기서 이분탐색을 활용할 수 있습니다. 이분탐색을 2번 사용해서 해당 숫자의 시작부분과 끝부분을 찾으면 구간 내 대표 숫자의 개수를 셀 수 있습니다. 물론, 이분탐색은 각각의 묶음 별로 진행해야 하며 만약 √N개씩 묶음으로 나누었다면 나눌 때 따로 정렬 하는 과정이 필요합니다. 

 

여기까지 3개 단계를 거치면 각각의 대표 숫자 별로 구간[a, b] 안에 몇 개가 있는지 알 수 있으며, 우리는 그 중에서 절반 이상의 개수를 가진 대표 숫자를 출력하면 됩니다. 자세한 소스코드는 다음과 같습니다.

 

소스코드

다음 소스코드는 머지소트트리를 사용한 소스코드입니다.

#include <stdio.h>
#include <vector>
#include <utility>
#define NMAX 300010
#define PAIR pair<int,int>
using namespace std;
 
int N, C, t, sz;
vector< int > mergeTree[NMAX*4];
int maxColor[NMAX*4];
 
PAIR tmp, vmax, ans;
 
int Q, a, b;
 
// 머지소트트리의 현재 인덱스(idx)에서 특정 색깔(val)의 개수 반환
int getColorCnt(int idx, int val) {
    int l, r, mid;
    int idxl, idxr;
 
    // 가장 왼쪽에 위치한 val의 위치 구하기
    l = 0; r = mergeTree[idx].size()-1;
    while(l<=r) {
        mid = (l+r)/2;
 
        if(mergeTree[idx][mid] < val) l = mid+1;
        else r = mid-1;
    }
    if(mergeTree[idx][l] != val) return 0;
    else idxl = l;
 
    // 가장 오른쪽에 위치한 val의 위치 구하기
    l = 0; r = mergeTree[idx].size()-1;
    while(l<=r) {
        mid = (l+r)/2;
 
        if(mergeTree[idx][mid] <= val) l = mid+1;
        else r = mid-1;
    }
    idxr = r;
 
    return (idxr-idxl)+1;
 
}
 
// 구간[l, r]에서 색깔(val)의 개수 가져오기
PAIR searchColor(int l, int r, int val) {
    PAIR ret;
 
    ret = make_pair(0, val);
    while(l<=r) {
        if(l%2 == 1) ret.first += getColorCnt(l++, val);
        if(r%2 == 0) ret.first += getColorCnt(r--, val);
 
        l/=2; r/=2;
    }
 
    return ret;
}
 
// 구간[l, r]에서 가장 많은 색깔과 개수 탐색
PAIR search(int l, int r) {
    int ll, rr;
    PAIR ret;
 
    ll = l; rr = r;
    ret = make_pair(0, 0);
    while(ll<=rr) {
        // maxColor[ll]/maxColor[rr]의 개수 중 최댓값 구하기
        if(ll%2 == 1) ret = max( ret, searchColor(l, r, maxColor[ll++]) );
        if(rr%2 == 0) ret = max( ret, searchColor(l, r, maxColor[rr--]) );
 
        ll/=2; rr/=2;
    }
 
    return ret;
}
 
 
// 머지소트트리의 현재 인덱스(idx)에서 가장 많은 색깔 찾기
void cntColor(int idx) {
    if(tmp.second == mergeTree[idx].back()) tmp.first++;
    else {
        vmax = max( vmax, tmp );
        tmp = make_pair(1, mergeTree[idx].back());
    }
}
 
int main() {
    // input
    scanf("%d %d", &N, &C);
    for(sz=1;sz<N;sz*=2);
    for(int i=0;i<N;i++) {
        scanf("%d", &t);
        // 바로 머지소트트리에 넣기
        mergeTree[sz+i].push_back(t);
 
        // 현재 구간에서 가장 많은 색깔 넣기
        maxColor[sz+i] = t;
    }
 
    // 머지소트트리 만들기
    for(int i=sz-1;i>0;i--) {
        int idxl, idxr, l, r;
 
        // 초깃값
        tmp = vmax = {0,0};
        l = i*2; r = i*2+1;
        idxl = idxr = 0;
 
        // 머지소트
        while(idxl<mergeTree[l].size() and idxr<mergeTree[r].size()) {
            if(mergeTree[l][idxl] < mergeTree[r][idxr]) mergeTree[i].push_back(mergeTree[l][idxl++]);
            else mergeTree[i].push_back(mergeTree[r][idxr++]);
 
            cntColor(i);
        }
 
        // 나머지 값 처리
        while(idxl<mergeTree[l].size()) {
            mergeTree[i].push_back(mergeTree[l][idxl++]);
            cntColor(i);
        }
        while(idxr<mergeTree[r].size()) {
            mergeTree[i].push_back(mergeTree[r][idxr++]);
            cntColor(i);
        }
 
        // 현재 구간에서 가장 많은 색
        vmax = max( vmax, tmp );
        maxColor[i] = vmax.second;
    }
 
 
    // 쿼리 처리
    scanf("%d", &Q);
    for(int i=1;i<=Q;i++) {
        scanf("%d %d", &a, &b);
        a--; b--;
 
        ans = search(a+sz, b+sz);
        if(ans.first > (b-a+1)/2) printf("yes %d\n", ans.second);
        else printf("no\n");
    }
}

후기

정답은 항상 대표 숫자 중 하나라는 조건을 떠올리기 쉽지 않은 문제였습니다. 아이디어에 비해 구현은 그리 어렵지 않았기에 쿼리 문제를 연습하는 사람에게 추천하고 싶은 문제입니다. mo's algorithm으로도 해결할 수 있으나 이부분은 추후에 다시 연습해볼 예정입니다.