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

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

3653번: 영화 수집

세그먼트 트리 연습 문제로 풀게 된 문제

1시간을 고민해도 아이디어가 전혀 안 떠올라서 블로그 3~4개 풀이를 다 읽어보고 겨우 이해했다.

그리고 구현하는데 또 시간초과 나서 구현 디테일 수정하느라 30분은 또 쓴 것 같다.

 

영화를 볼 때마다 꺼낸 영화보다 위에 있던 영화들의 위치가 한칸씩 밀리는 것이 이 문제의 어려운 점이다.

이는 N+M개의 리프노드를 가지는 세그트리를 구현하여 해결할 수 있다.

 

문제의 첫번째 테스트 케이스를 이미지로 설명하면 아래와 같다.

위 이미지와 같이 N+M 사이즈의 배열을 구성한다. (왼쪽 숫자는 인덱스이다.)

N개의 노드는 미리 1, 2, 3 영화로 채워둔다.

 

1. 영화 위치를 찾는다.

3번 영화의 위치를 찾는다.

3번 영화의 위치는 위치를 기록하는 별도 배열을 통해 찾는다.

3번 영화의 위치가 0번 인덱스 임을 알았다.

 

2. 그 영화 위에 있는 영화 개수를 구한다.

3번 인덱스 위치의 다음 인덱스인 1번 인덱스부터 마지막 5번 인덱스까지 영화가 들어있는 칸의 개수 (구간합)를 구한다. 1, 2번 인덱스 칸에만 들어있으니 2가 답이다.

 

3. 찾은 영화를 위로 올린다.

기존 위치인 0번 인덱스를 비우고, 이를 구간합에 반영한다.

3번 영화를 아직 방문한 적 없는 인덱스 중 가장 낮은 값인 3번 인덱스로 올린다.

마찬가지로 이를 구간합에 반영한다.

---

모든 쿼리에 대해 이를 반복하면 된다.

다음으로 1번 영화를 찾는다.

 

1. 영화 위치를 찾는다.

1번 영화의 위치를 찾는다.

1번 영화의 위치는 위치를 기록하는 별도 배열을 통해 찾는다.

1번 영화의 위치가 2번 인덱스 임을 알았다.

 

2. 그 영화 위에 있는 영화 개수를 구한다.

1번 인덱스 위치의 다음 인덱스인 3번 인덱스부터 마지막 5번 인덱스까지 영화가 들어있는 칸의 개수 (구간합)를 구한다.

3번 인덱스 칸에만 들어있으니 1이 답이다.

 

3. 찾은 영화를 위로 올린다.

기존 위치인 2번 인덱스를 비우고, 이를 구간합에 반영한다.

1번 영화를 아직 방문한 적 없는 인덱스 중 가장 낮은 값인 4번 인덱스로 올린다.

마찬가지로 이를 구간합에 반영한다.

---

다시 1번 영화를 찾는다.

 

1. 영화 위치를 찾는다.

1번 영화의 위치를 찾는다.

1번 영화의 위치는 위치를 기록하는 별도 배열을 통해 찾는다.

1번 영화의 위치가 4번 인덱스 임을 알았다.

 

2. 그 영화 위에 있는 영화 개수를 구한다.

4번 인덱스 위치의 다음 인덱스인 5번 인덱스부터 마지막 5번 인덱스까지 영화가 들어있는 칸의 개수 (구간합)를 구한다.

모든 인덱스가 비어있으므로 0이 답이다.

 

3. 찾은 영화를 위로 올린다.

기존 위치인 4번 인덱스를 비우고, 이를 구간합에 반영한다.

1번 영화를 아직 방문한 적 없는 인덱스 중 가장 낮은 값인 5번 인덱스로 올린다.

마찬가지로 이를 구간합에 반영한다.

 

따라서 2 1 0 이 출력으로 나온다.

---

이제 문제를 푸는 방법을 알았으니 구현하면 된다.

나는 처음에 기존 구간합 구하기 문제처럼 구현했다가 시간초과를 받았다.

업데이트를 할 때 조금 더 빠르게 업데이트를 할 수 있어야 한다.

 

def update(node, start, end, prev_pos, now_pos):
    if start == end:
        if start == prev_pos:
            seg_tree[node] = 0
        elif start == now_pos:
            seg_tree[node] = 1
        return seg_tree[node]

    mid = (start + end) // 2
    left_tree = update(node*2, start, mid, prev_pos, now_pos)
    right_tree = update(node*2+1, mid+1, end, prev_pos, now_pos)
    seg_tree[node] = left_tree + right_tree
    return seg_tree[node]

 

내가 처음에 작성했던 업데이트 함수이다.

값은 반드시 1->0 또는 0->1 로 바뀌기만 하는데, 값에 변화가 생긴 이후, 그 인덱스가 포함된 범위를 나타내는 모든 노드의 값을 다시 더해서 계산하고 있다.

이 부분이 시간초과로 작용하였다.

 

큰 수를 연산하고, 이를 return 해주고, return 받은 값을 다시 더해주고...

이 부분에서 리소스를 많이 사용한 것 같다.

def update(node, start, end, pos, diff):
    if pos < start or end < pos:
        return

    seg_tree[node] += diff

    if start == end:
        return

    mid = (start + end) // 2
    update(node*2, start, mid, pos, diff)
    update(node*2+1, mid+1, end, pos, diff)

 

이 코드는 다른 블로그의 코드를 보고 개선한 업데이트 함수이다.

특정 점의 값이 변하면, 그 점이 포함된 모든 구간 노드에 대해 그 변화량만 적용해준다.

연산하는 숫자의 크기가 매우 작아졌으며, return 값이 필요하지 않아 속도가 훨씬 올라간다.

 

 

아주 아슬아슬하게 맞을 수 있었다.

 

다른 정답코드 중에서, 세그먼트 트리를 재귀가 아니라 반복문으로 구현한 경우도 보였는데, 이는 아직 이해를 하지 못하였다...

또 이 문제를 푸는 유형의 세그트리를 펜윅트리 / BIT 등으로 명명하여 풀기도 하는데, 이 부분도 아직 이해를 하지 못했다.

일단 기본 세그트리부터 익숙해져봐야겠다.

 

import sys
input = sys.stdin.readline

def init(node, start, end):
    global n
    if start == end:
        if start <= n:
            seg_tree[node] = 1
        return seg_tree[node]

    mid = (start + end) // 2
    left = init(node*2, start, mid)
    right = init(node*2+1, mid+1, end)

    seg_tree[node] = left + right
    return seg_tree[node]


def query(node, start, end, left, right):
    # left ~ right 가 구하려는 구간
    if end < left or right < start: # 현재 보고 있는 구간이 구하는 구간 밖의 범위
        return 0

    if left <= start and end <= right: # 현재 보고 있는 구간이 구하는 구간에 포함
        return seg_tree[node]

    mid = (start + end) // 2
    left_tree = query(node*2, start, mid, left, right)
    right_tree = query(node*2 + 1, mid+1, end, left, right)
    return left_tree + right_tree


def update(node, start, end, pos, diff):
    if pos < start or end < pos:
        return

    seg_tree[node] += diff

    if start == end:
        return

    mid = (start + end) // 2
    update(node*2, start, mid, pos, diff)
    update(node*2+1, mid+1, end, pos, diff)


T = int(input())
for _ in range(T):
    n, m = map(int, input().split())
    queries = map(int, input().split())
    position = [0] + [i for i in range(n, -1, -1)]
    seg_tree = [0 for _ in range((n+m)*4)]

    init(1, 1, n+m)

    query_count = 0
    for movie in queries:
        pos = position[movie]
        query_count += 1
        position[movie] = n+query_count
        print(query(1, 1, n+m, pos+1, n+m), end=' ')
        # update(1, 1, n+m, pos, position[movie])
        update(1, 1, n+m, pos, -1)
        #print(seg_tree)
        update(1, 1, n+m, position[movie], 1)
    print()