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) 등 여러 문제에서 접해보았기 때문에 쉽게 해결할 수 있는 문제였다.
오늘의 교훈) 팬윅트리는 매우 유용하다