수열과 쿼리 1 [BOJ 13537]

수쿼 시리즈의 화려한 도약을 알리는 매우 재미있는 문제이다. 머지 소트 트리 기본 문제이면서도, 이 문제는 머지 소트 트리 외에 별해가 존재한다. 개인적으로 별해의 아이디어가 상당히 의미 있다고 생각한다. 아래 풀이에 별해와 관련된 내용은 모두 가려놓았으니, 최소 30분 정도는 별해에 대해 고민해보는 것을 추천한다.

 

** 수열과 쿼리 3 풀다가 오신 분들은 머지 소트 트리에 대한 풀이인 풀이 1만 읽으면 수열과 쿼리 3을 풀 수 있다.

 

문제 요약

길이 10만 이하의 수열 A에서 다음 쿼리를 수행한다.

i j k : 구간 [i, j]에서 k보다 큰 원소의 수를 출력한다.

 

사용 알고리즘

풀이 1) 머지 소트 트리

풀이 2) Hint : 오프라인 쿼리라는 점을 이용

스위핑 + 세그먼트 트리

 

풀이

풀이 1)

머지 소트 트리의 가장 기본적인 문제이고, 그대로 구현해주면 AC를 받는다.

필자는 구현 시 std::merge 함수를 사용했다. 정렬된 2개의 vector를 하나로 merge해주는 함수로, 선형의 시간복잡도에 돌아간다. 굳이 저런 것까지 알아야하나 싶긴한데, 개인적으로 머지 소트 트리를 구현할 때 매우 편리하다고 생각하여 활용법을 익혀두었다. 실제로 필자의 코드는 1000B를 넘지 않았다. merge 함수를 사용하기 전 vector를 크기에 맞게 resize하는 것을 잊지 말자.

시간복잡도는 O(NlogN+Q(logN)^2)이다. 로그 제곱 스케일이라 조금 오래 걸린다.

풀이 2)

오프라인 쿼리는 일반적으로 "쿼리 정렬"에 대해 다룬다. 하지만 이 문제에서는 여기서 그치지 않고, 한 발 더 나아가 "수열 정렬"을 적용함으로써 풀 수 있다.

우선 수열을 내림차순으로 정렬하고, 쿼리는 k를 기준으로 내림차순으로 정렬한다. 그리고 세그먼트 트리를 하나 잡는데, 이 세그먼트 트리에서 구간 [i, j]를 관리하는 노드는 구간 [i, j]에서 k보다 큰 원소의 수를 저장한다. 쿼리가 진행될수록 k가 작아지고, 이에 따라 수열의 각 원소들은 스위핑을 통해 최대 1번 세그먼트 트리에 추가된다. 각 쿼리의 출력값은 세그먼트 트리에 쿼리를 하나 날려서 풀 수 있다. 사고의 과정이 조금 복잡한 거 같지만, 이 풀이를 사용한 필자의 코드도 1000B를 넘지 않았다. 직접 구현한다면 확실히 이해하게 될 것이다. 필자는 상수를 줄이기 위해 펜윅 트리로 구현하였다.

시간복잡도는 O((Q+N)logN+MlogM)이다. 시간복잡도에서 로그가 하나 떼어졌다.

아래 제출이 풀이 1, 위 제출이 풀이 2에 대한 코드이다.

실제 시간도 256ms에서 76ms로 약 30%가 되었고, 메모리 사용량은 25% 정도로 줄었다. 머지 소트 트리는 4N개의 노드가 벡터 하나씩을 저장하므로 메모리를 매우 많이 차지하게 되지만, 펜윅 트리는 N개의 노드가 정수 하나를 저장할 뿐이기에 메모리를 거의 먹지 않는다. 이 문제가 아니더라도 특정 아이디어를 통해 머지 소트 트리 문제를 세그먼트 트리 문제로 만들어 시간 및 메모리를 크게 줄일 수 있다. 제한이 빡빡한 문제를 만나면 써먹을 수 있도록 이 문제에서 연습해두자.

 

코드

풀이 1에 대한 코드)

#include<bits/stdc++.h>
#define SIZE 100010
using namespace std;
vector<int>seg[SIZE<<2];
int N,M,A[SIZE],i,j,k;
void init(int l,int r,int node)
{
    if(l==r){
        seg[node].push_back(A[l]);
        return;
    }
    int mid=l+r>>1;
    init(l,mid,node<<1);
    init(mid+1,r,(node<<1)+1);
    seg[node].resize(seg[node<<1].size()+seg[(node<<1)+1].size());
    merge(seg[node<<1].begin(),seg[node<<1].end(),seg[(node<<1)+1].begin(),seg[(node<<1)+1].end(),seg[node].begin());
}
int query(int l,int r,int node,int s,int e,int k)
{
    if(l>e||r<s)return 0;
    if(s<=l&&r<=e)return seg[node].end()-lower_bound(seg[node].begin(),seg[node].end(),k+1);
    int mid=l+r>>1;
    return query(l,mid,node<<1,s,e,k)+query(mid+1,r,(node<<1)+1,s,e,k);
}
int main()
{
    ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>N;
    for(i=1;i<=N;i++)cin>>A[i];
    init(1,N,1);
    for(cin>>M;M--;){
        cin>>i>>j>>k;
        cout<<query(1,N,1,i,j,k)<<'\n';
    }
}

풀이 2에 대한 코드)

#include<bits/stdc++.h>
#define SIZE 100010
using namespace std;
struct Query
{
    int x,y,z,idx;
};
int N,M,seg[SIZE],ans[SIZE],c=1;
pair<int,int>A[SIZE];
Query Q[SIZE];
bool comp(pair<int,int>a,pair<int,int>b)
{
    return a.first>b.first;
}
bool comp2(Query a,Query b)
{
    return a.z>b.z;
}
void add(int s)
{
    while(s<=N)seg[s]++,s+=(s&-s);
}
int get(int s)
{
    int ret=0;
    while(s>0)ret+=seg[s],s-=(s&-s);
    return ret;
}
int main()
{
    ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>N;
    for(int i=1;i<=N;i++){
        cin>>A[i].first;
        A[i].second=i;
    }
    cin>>M;
    for(int i=1;i<=M;i++){
        cin>>Q[i].x>>Q[i].y>>Q[i].z;
        Q[i].idx=i;
    }
    sort(A+1,A+N+1,comp);
    sort(Q+1,Q+M+1,comp2);
    for(int i=1;i<=M;i++){
        for(;c<=N;c++){
            if(A[c].first>Q[i].z)add(A[c].second);
            else break;
        }
        ans[Q[i].idx]=get(Q[i].y)-get(Q[i].x-1);
    }
    for(int i=1;i<=M;i++)cout<<ans[i]<<'\n';
}