수열과 쿼리 9 [BOJ 13554]

이 문제는 내가 N달 동안 고민한 문제로, 자구충이라면 다4 수쿼 정도는 풀 수 있어야하지 않을까 싶어서 업솔빙도 안했다. 그리고 최근에 다시 고민을 하다가 문제를 푸는데에 핵심적인 아이디어를 떠올려 풀 수 있게 되었다. vector의 몇 안되는 다이아 이상 자력솔 중 하나이다.

 

문제 요약

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

i j k : i<=p,q<=j이면서 A[p]*B[q]<=k인 (p, q) 쌍의 개수를 출력한다.

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

 

사용 알고리즘

Mo's Algorithm

펜윅 트리

 

풀이

 

관찰 1 : 업데이트가 없는 오프라인 쿼리 + 시간제한 6초

오프라인 쿼리란 업데이트 쿼리가 없어 쿼리의 순서를 바꾸어도 무방한 쿼리를 말한다.

오프라인 쿼리를 적당히 정렬하여 쿼리 간 이동을 O(NsqrtN)에 할 수 있는 Mo's Algorithm이 자연스럽게 생각난다.

마침 시간제한도 6초로 N의 제한인 10만에 비해 상당히 길어 시간복잡도에 sqrtN이 들어가는 것을 추측해줄 수 있다.

 

관찰 2 : 변하는 k

Mo's Algorithm을 사용하는 가장 기초적인 문제들에서는 주로 쿼리 간 이동을 하며 이전 쿼리의 출력값을 통해 쿼리의 출력값을 즉시 계산한다. 즉, 쿼리 당 시간복잡도가 O(1)이 되어 전체 시간복잡도가 O(NsqrtN)이 되는 경우가 많다.

하지만 이 경우는 조금 다르다. 우리가 어떤 구간 [i1, j1]에 대해 i<=p, q<=j이면서 A[p]*B[q]<=k1인 것의 개수를 알고 있다고 해보자. 이 경우 다음 쿼리의 k값이 변한다면, 현재 쿼리의 답은 무용지물이 된다.

따라서 각 쿼리의 답을 구할 때 가장 많은 영향을 미치는 값인 k에 대한 쿼리를 날려줌으로써 문제를 해결해야할 것 같다는 느낌이 든다. 여기서 세그먼트 트리를 떠올리는 것은 어렵지 않다. A[p]*B[q]<=k인 B[q]의 개수를 알기 위해서는 k/A[p]이하의 B[q] 개수를 알면 되기 때문에, 쿼리 간 이동에서 값들이 추가/삭제되는 것을 세그먼트 트리에 업데이트하면 적당히 해결할 수 있을 것 같다. 수열 값들의 범위가 모두 10만 이하여서 좌표 압축은 필요하지 않다.

 

관찰 3 : 쿼리 횟수 줄이기

하지만 현재 구간 내에 모든 A[i]에 대해 쿼리를 날린다면, 혹은 1부터 k까지 모든 값에 대해 쿼리를 날린다면 시간복잡도가 보장될 수 없다. 대부분 관찰 2까지는 쉽게 하는데, 관찰 3을 하지 못해 이 문제를 못 푸는 경우가 많다. 나도 관찰 3을 못해서 N달을 고민하였고, 서론에서 찾은 핵심적인 아이디어가 바로 관찰 3이다.

어떤 두 수를 곱해 k가 된다고 하면, 두 수 중 하나가 sqrt(k) 이하임이 자명하다. 이를 이용하면 각 쿼리 당 세그먼트 트리 쿼리를 최대 2sqrt(k)개까지만 날려주면 된다. 단, 두 수가 모두 sqrt(k) 이하인 경우가 2번 세어지지 않도록 주의한다.

 

이렇게 해서 풀이가 완성되었다. 총 시간복잡도는 O(NsqrtN+QsqrtKlogN)이다. 약간 아슬아슬해 보이는데, 일반 세그먼트 트리는 상수가 커서 시간 초과가 난다. 펜윅 트리가 가능한 꼴이므로 펜윅 트리로 구현해주면 AC를 받는다. 필자의 코드는 2776ms가 나왔다. (만약 시간복잡도가 보장되는 상태에서 세그먼트 트리가 시간초과를 받는다면 본인이 시간복잡도를 잘못 계산했거나 잘못 구현했을 확률이 제일 높지만 펜윅 트리가 가능한 경우 펜윅 트리를, 불가능한 경우 바텀업 세그먼트 트리를 구현해보면 상수를 줄일 수 있다. 세그->바텀업 세그로 AC를 받은 경우는 그렇게 많이 보지는 못했는데, 세그->펜윅으로 AC를 받은 경우는 많이 봐서 펜윅은 꼭 익혀두길 바란다.)

 

코드

#include<bits/stdc++.h>
#define SIZE 100010
using namespace std;
typedef long long ll;
struct Queryst
{
    int S,E,K,I;
};
Queryst Q[SIZE];
int N,M,A[2][SIZE],sqrtN;
ll seg[2][SIZE],ans[SIZE];
bool comp(Queryst x,Queryst y)
{
    if(x.S/sqrtN==y.S/sqrtN)return x.E<y.E;
    return x.S<y.S;
}
void update(int s,ll w,int idx)
{
    s=A[idx][s];
    while(s<=SIZE){
        seg[idx][s]+=w;
        s+=(s&-s);
    }
}
ll sum(int s,int idx)
{
    ll ret=0;
    while(s>0){
        ret+=seg[idx][s];
        s-=(s&-s);
    }
    return ret;
}
ll query(int s,int e,int idx)
{
    return sum(e,idx)-sum(s-1,idx);
}
int main()
{
    cin>>N;sqrtN=sqrt(N);
    for(int i=0;i<2;i++)for(int j=1;j<=N;j++)scanf("%d",&A[i][j]);
    cin>>M;
    for(int i=0;i<M;i++){
        scanf("%d %d %d",&Q[i].S,&Q[i].E,&Q[i].K);
        Q[i].I=i;
    }
    sort(Q,Q+M,comp);
    int s=Q[0].S,e=Q[0].E,k=Q[0].K,idx=Q[0].I;
    for(int i=s;i<=e;i++){
        update(i,1,0);
        update(i,1,1);
    }
    for(int i=1;i<=sqrt(k);i++){
    	ans[idx]+=query(i,i,0)*query(1,k/i,1)+query(i,i,1)*query(sqrt(k)+1,k/i,0);
    }
    for(int i=1;i<M;i++){
        k=Q[i].K,idx=Q[i].I;
        while(s<Q[i].S){
            update(s,-1,0);
            update(s,-1,1);
            s++;
        }
        while(s>Q[i].S){
            s--;
            update(s,1,0);
            update(s,1,1);
        }
        while(e<Q[i].E){
            e++;
            update(e,1,0);
            update(e,1,1);
        }
        while(e>Q[i].E){
            update(e,-1,0);
            update(e,-1,1);
            e--;
        }
        for(int i=1;i<=sqrt(k);i++){
        	ans[idx]+=query(i,i,0)*query(1,k/i,1)+query(i,i,1)*query(sqrt(k)+1,k/i,0);
        }
    }
    for(int i=0;i<M;i++)printf("%lld\n",ans[i]);
}