BOJ 1168 요세푸스 문제 2

 

문제 링크 : https://www.acmicpc.net/problem/1168

1168번: 요세푸스 문제 2

 

0 으로 시작하는 index 를 다루면서 개수를 따지는 문제는 언제나 헷갈린다.(언제쯤이면 실수 없이 쓸 수 있을까?)

 

문제에서 1 부터 N 까지의 수를 다루지만, 원형으로 배치된 상태에서 진행되는 상황은 mod 연산이 필수적이다.

mod 연산에서 헷갈리지 않게 1 부터 N 까지를 0 부터 N-1 로 바꾸어 생각하자. 

여기까지는 어렵지 않다.

 

문제는 'K 번째' 에 있다.  

아래 코드에서 사용되는 qry  함수의 마지막 인자 k 는 '개수'이다.  

 

 

qry 는 찾고자 하는 개수가 왼쪽 자식이 보유한 것보다 작거나 같으면 왼쪽 자식에게,

그렇지 않으며 오른쪽 자식에게 질문을 넘긴다.

 

반면, 우리는 0-base index를 사용해야 하므로 k 개가 되는 index 는 k-1 이다.

아래 코드에서 for 루프 직전의 

 

 

에서 이를 반영했다. 즉, index 와 개수 사이의 차이가 1을 유지해야 한다.

또한 원소가 하나 제거될 때마다 sz 를 1 씩 줄여가면서 mod 연산을 수행해야 한다. 

 

기껏해야 덧셈에서도 지금 더한 것이 index 위의 연산인지 개수 위에서 연산인지를 구분하여 사고하는 것은 너무 힘들다.

 

#include<bits/stdc++.h>

using namespace std;

int sz, k;
int tree[400000];

void init(int l,int r,int node){
    if(l==r) {
        tree[node]=1;
        return;
    }
    
    int mid=l+r>>1;
    init(l,mid,2*node);
    init(mid+1,r,2*node+1);
    tree[node]=tree[2*node]+tree[2*node+1];
}

void upd(int l, int r, int node, int p, int v=-1){
    if(p<l || r<p) return;
    if(l==r) {
        tree[node]+=v;
        return;
    }
    
    int mid=l+r>>1;
    upd(l,mid,2*node,p,v);
    upd(mid+1,r,2*node+1,p,v);
    tree[node]=tree[2*node]+tree[2*node+1];
}

int qry(int l, int r, int node, int k){
    if(l==r) return l;
    
    int mid=l+r>>1;
    int d=tree[2*node];
    if(d>=k) return qry(l,mid,2*node, k);
    else return qry(mid+1,r,2*node+1, k-d);
}

int main(){
    int N,K; scanf("%d%d",&N,&K);
    
    init(0, N-1, 1);
    
    printf("<");
    sz=N, k=K-1;
    for(int i=0;i<N-1;++i){
        int p=qry(0, N-1, 1, k+1);
        printf("%d, ", p+1);
        upd(0, N-1, 1, p);
        sz--;
        k=(k-1+K)%sz;
    }
    printf("%d>\n", qry(0,N-1, 1, 1)+1);
}

 

위 코드에서 k 는 찾아야 하는 index를 뜻한다고 말했다. 

만약 k 를 개수의 의미로 두고 싶다면  다음과 같이 해도 된다.

 

    sz=N, k=K;
    for(int i=0;i<N-1;++i){
        int p=qry(0, N-1, 1, k);
        printf("%d, ", p+1);
        upd(0, N-1, 1, p);
        sz--;
        k=(k-2+K)%sz+1;
    }

 

이 코드에서 등장하는 

k=(k-2+K)%sz+1;

는 내가 작성하고서도 매번 헷갈리기 때문에 앞선 코드를 선호한다.