[백준 1168번] 요세푸스 문제 2 (Python3)

def update(i,v):
  while i<=M:
    fen[i] += v
    i += i&-i

def cal(i):
  s = 0
  while i:
    s += fen[i]
    i -= i&-i
  return s

def find(x,i,c):
  s = cal(i)
  if s<x:
    return find(x,i+(1<<c),c-1)
  if s>x or s==x and check[i] or i>N:
    return find(x,i-(1<<c),c-1)
  return i

N,K = map(int,input().split()); M = 1<<17
fen = [0]*(M+1); check = [0]*(M+1)
for i in range(1,N+1):
  update(i,1)
seq = [0]
for _ in range(N):
  x = (cal(seq[-1])+K)%fen[M]
  if not x:
    x = fen[M]
  n = find(x,M,16)
  update(n,-1); check[n] = 1; seq.append(n)
print("<"+", ".join(map(str,seq[1:]))+">")

팬윅트리로 해결하였다.

처음 이 문제를 봤을 때는 풀이가 전혀 떠오르지 않던 문제였다. 그런데 세그먼트 트리 (혹은 팬윅트리) 활용 문제를 여러개 풀고 나니 쉽게 해결할 수 있었다.

팬윅트리나 세그먼트 트리를 이용해서 K번째 원소를 찾는 방법을 알고 있다면 쉽게 해결할 수 있는 문제

 

풀이를 설명하면,

1. 팬윅트리로 1~N 인덱스에 1씩 갱신한다.

2. 가장 처음은 0부터 시작해서 현재값+K인 인덱스를 팬윅트리+이분탐색을 통해 찾아준다. 찾았다면 check하고, 팬윅트리는 해당 인덱스에 -1을 갱신한다.

3. 만약 현재값+K가 전체합보다 더 큰 경우, 전체합에 대한 나머지만큼을 찾아주어야 한다. (예를들어 7번째 수를 찾아야 하는데 현재 남은 숫자가 3개라면 7%3=1, 1번째 수를 찾는다)

4. 구한 요세푸스 순열을 출력형식에 맞게 출력한다.

 

각 값에 1씩 갱신해두고 이분탐색을 통해 K번째 수를 찾는 것은 영화 수집 [백준 3653번] 영화 수집 (Python3) (tistory.com) , 이진탐색트리 [백준 2957번] 이진 탐색 트리 (Python3) (tistory.com) 등 여러 문제에서 접해보았기 때문에 쉽게 해결할 수 있는 문제였다.

 

 

오늘의 교훈) 팬윅트리는 매우 유용하다