(1168) 요세푸스 문제 2

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

단계별로 풀어보기

44단계(세그먼트 트리 1) 8번째

 

 

 

이전에 풀었던 요세푸스 문제의 범위가 늘어나 큐로 일일이 시뮬레이션하면 O(NK)로 시간초과가 난다.

이전 문제와 같은 방식으로 1~N번 배열의 값을 1로 만들고, 현재 살아있는 사람 중 K번째 사람을 찾는다. 찾으면 출력하고 그 사람은 제외되었으므로 인덱스에 해당하는 값을 0으로 만든다. 다음으로 제외될 K'번째 사람은 남은 사람 수나 K 등의 값, 그리고 모듈러 연산을 이용해 수학적으로 구할 수 있다. 새그먼트 트리를 이용해 시간복잡도가 O(N log N)으로 줄어들었다.

 

 

 

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  int n, k;
  
  cin >> n >> k;
  
  int treesize= (1<<(((int)ceil(log2(n+1)))+1));
  vector<long> segtree(treesize);
  
  auto init= [&](auto self, int s, int e, int node) -> void {
    if(s==e){
      segtree[node]= 1;
    }else{
      self(self, s, s+(e-s)/2, node*2);
      self(self, s+(e-s)/2+1, e, node*2+1);
      segtree[node]= segtree[node*2]+segtree[node*2+1];
    }
  };
  
  auto update= [&](auto self, int s, int e, int node, int idx) -> void {
    if(idx<s||e<idx)  return;
    
    if(s==e){
      segtree[node]= 0;
    }else{
      self(self, s, s+(e-s)/2, node*2, idx);
      self(self, s+(e-s)/2+1, e, node*2+1, idx);
      segtree[node]= segtree[node*2]+segtree[node*2+1];
    }
  };
  
  auto query= [&](auto self, int s, int e, int node, int x) -> int {
    if(s==e)  return s;
    
    int lcnt= segtree[node*2];
    
    if(x<=lcnt)  return self(self, s, s+(e-s)/2, node*2, x);
    else  return self(self, s+(e-s)/2+1, e, node*2+1, x-lcnt);
  };
  
  init(init, 1, n, 1);
  
  cout << "<";
  for(int i=n-1,x=k;i>=0;i--){
    int res= query(query, 1, n, 1, x);
    
    cout << res << (i?", ":"");
    
    update(update, 1, n, 1, res);
    
    x+= k-1;
    if(!i)  x= 1;
    else  x= x%i ? x%i : i;
  }
  cout << ">";
  
  return 0;
}