[백준 3653번] 영화 수집 (Python3)

import sys
input = sys.stdin.readline

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

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

for _ in range(int(input())):
  N,M = map(int,input().split()); K = 1<<18
  movie = [*map(int,input().split())]
  
  Index = [i+(1<<17) for i in range(N+1)]; top = 1<<17
  fen = [0]*K
  
  for n in range(1,N+1):
    update(Index[n],1)
  
  result = []
  for n in movie:
    result.append(SUM(Index[n]-1))
    update(Index[n],-1)
    Index[n] = top; update(top,1); top -= 1
  print(*result)

팬윅트리로 해결하였다.

아이디어를 떠올리기 어려운 문제였다. 꽤 오랫동안 풀지 못하다가 저장위치의 개수를 N+M개로 해야한다는 힌트를 듣고 풀 수 있었다.

 

풀이를 설명하면,

1. 가장 처음에는 모든 책을 i+2^17의 위치에 두었다고 가정하고 위치를 저장한다. 그리고 책이 없는 가장 위를 top = 2^17로 저장한다.

2. 팬윅트리에 모든 책의 위치에 1씩 갱신한다.

3. 찾고자 하는 책의 위치의 앞쪽에 책이 몇권 있는지를 결과에 저장하고, 팬윅트리에 해당 위치에 -1을 해준 뒤 top 위치로 책을 옮기고 그 위치에 팬윅트리를 갱신한 뒤 top을 -1 해준다.

4. 3에서 저장한 결과들을 출력한다.

 

저장위치의 개수를 N개로 해야한다고 속단하고 문제를 해결하려 하면 풀기 어려운 문제였다. 저장위치를 N+M, 혹은 넉넉하게 2^18개로 하는 것이 핵심 포인트였다.

 

 

오늘의 교훈) 위치를 옮길때 저장공간의 크기를 어떻게 할지 고민하자