[백준] 3653 - 영화 수집 (Python)

sudog·2023년 8월 27일
0

백준

목록 보기
14/15

DVD를 꺼내서 본 후에 가장 위에 놓는 배열을 구현하는 것은 어렵지 않다. 하지만 단순한 배열로 구현하면 각 DVD의 위에 놓인 영화의 개수를 찾는 것이 비효율적인 연산이 될 것이다.

따라서 구간의 DVD개수를 빠르게 구할 수 있는 세그먼트 트리를 사용해야 한다. 이 문제에서는 누적 개수를 세면 해결이 가능하므로 더 효율적인 펜윅 트리를 사용해보자.

import sys
input = sys.stdin.readline

t = int(input())

def init_tree():
	# 영화를 m번 봐야 하므로 펜윅 트리의 크기는 n+m+1이 된다.
    for i in range(1, n+m+1):
    	# 초기에 인덱스 n까지는 모두 DVD가 차 있으므로 구간의 길이가 곧 DVD의 개수가 된다.
        if i <= n:
            fw_tree[i] = i & -i
        # n보다 큰 인덱스에는 DVD가 존재하지 않으므로 n보다 작거나 같은 구간만 더해준다.
        else:
            fw_tree[i] = max(0, n - (i - (i & -i)))
            
def update_tree(cur_idx, next_idx):
    temp = cur_idx
    result = 0
    while temp > 0:
        result += fw_tree[temp]
        temp -= temp & -temp
    while cur_idx <= n+m:
        fw_tree[cur_idx] -= 1
        cur_idx += cur_idx & -cur_idx
    while next_idx <= n+m:
        fw_tree[next_idx] += 1
        next_idx += next_idx & -next_idx
    return n - result
    
for _ in range(t):
    n, m = map(int, input().rstrip().split())
    fw_tree = [0]*(n+m+1)
    # 영화의 번호와 인덱스를 매핑(초기상태는 1이 가장 위에 존재)
    index_map = [n-i+1 for i in range(n+1)]
    init_tree()
    for i, e in enumerate(map(int, input().rstrip().split())):
        idx = index_map[e]
        # 가장 위에 존재하는 DVD를 꺼낸 경우 아무 변화도 없다.
        if idx == n+i+1:
            print(0, end=' ')
        else:
            print(update_tree(idx, n+i+1), end=' ')
            # 본 영화는 가장 위로 이동한다.
            index_map[e] = n+i+1
    print()
profile
안녕하세요

0개의 댓글