[백준/Python]23326: 홍익 투어리스트

문제

https://www.acmicpc.net/problem/23326
백준 - 문제집 - 0x16강 - 이진 검색 트리

 

알고리즘: 이진 검색 트리, 트리를 사용한 집합과 맵

---

풀이

21944: 문제 추천 시스템 Version 2와 마찬가지로 무한 TLE♾️에 빠지게 했던 녀석이다.

 

1️⃣ 첫번째 시도
명소들을 현재 위치: [현재 위치와 명소의 거리들] 형태의 딕셔너리로 구성.
각각의 값 리스트는 최소힙으로 사용함.
=> 어림도 없이 시간초과

 

2️⃣ 두번째 시도
이진탐색 모듈로 bisect_left, insort_left 함수를 사용.
명소들만 따로 리스트에 저장 후, 이진탐색으로 알맞은 인덱스를 찾아 삽입/삭제 + 탐색을 진행했음.
=> 66% 정도까지 진행 후 시간초과

 

Python으로는 안되는건가? 했지만 PyPy3로 통과된 분들이 존재했다.
맞혔다/틀렸다의 문제가 아니라 너무 답답했고... 결국 편법을 써서 해답 코드를 엿볼 수 있었다.
(다른 언어로 제출해서 통과한 후 공개된 풀이 코드 탐색^^;;)

 

🗝️ 정답 키 포인트

해결 방안은 세그먼트 트리(Segment Tree)를 사용하는것이었다.
사실 나도 세그먼트 트리를 이 문제로 처음 알게 되었다.

 

세그먼트 트리란?
: 특정 범위의 값을 빠르게 계산할 수 있도록 설계된 트리 자료구조

 

범위를 절반씩 나누어 범위의 값을 저장해두는 형태다.
원하는값보다 현재 세그먼트 트리 노드의 값이 작거나 같다면 왼쪽, 크다면 오른쪽으로 이동한다.
세그먼트 트리의 크기는 보통 4 * 트리의 크기 로 설정하는 편이다.

본인 마음대로지만, 루트노드를 1로 잡고 시작해야 왼쪽 자식/오른쪽 자식을 구분하기 쉽다.
1부터 시작할경우 현재 노드 번호를 node라고 할 때 왼쪽 자식의 번호는 2 * node, 오른쪽 자식의 번호는 2 * node + 1가 되므로 쉽게 선택할 수 있기 때문이다.

 

어쨌든 이진 탐색 트리이므로 원하는 타겟값과 일치할때까지 start + end의 중간값 mid와 타겟값을 비교해나간다.
왼쪽 서브트리로 이동 시 start는 그대로, end는 mid로 업데이트. (작거나 같은 경우에 해당되므로)
오른쪽 서브트리로 이동 시 start는 mid + 1로, end는 그대로 유지한다.

 

만약 명소들의 상태가 다음과 같다고 가정해보자.
[1, 0, 0, 1, 1, 0, 1, 0] = 1: 명소, 0: 안명소
0번째 인덱스부터 현재 위치까지의 명소의 갯수가 3개인데, 현재 위치가 어디인지 알고싶다면?

segt = [0] * 4 * 8  # 원래 트리크기의 4배

segt[1] = 4  # 전체범위의 명소 갯수

node = 1
start = 0
end = 7 (N-1)

mid = (start + end) // 2  # 7 // 2 = 3

# 왼쪽 서브트리 ( 2 * node )
segt[2] = 2  # 0~3 의 명소 갯수
# 오른쪽 서브트리 ( 2 * node + 1 )
segt[3] = 4  # 4~7 의 명소 갯수

# 찾고자 하는 값은 3이므로 오른쪽 서브트리로 이동해야함.
# 왼쪽범위의 값(2)를 뺀 후 오른쪽 서브트리로 이동.
# 찾고자 하는 값은 3 - 2 = 1. 이제 1이 됨.

node = 3
start = mid + 1  # 4
end = 7

mid = (4 + 7) // 2 = 5

# 오른쪽 서브트리의 서브트리
# 왼쪽 서브트리 ( 2 * node )
segt[6] = 1  # 4~5 의 명소 갯수
# 오른쪽 서브트리 ( 2 * node + 1 )
segt[7] = 1  # 6~7 의 명소 갯수

# segt[6] 값과 찾고자하는 값이 같으므로 왼쪽 서브트리로 이동
# 왼쪽 서브트리로 이동하므로 찾고자 하는 값은 그대로 유지

node = 6
start = 4
end = 5

mid = (4 + 5) // 2 = 4

# 왼쪽 서브트리
segt[12] = 1  # 4~4의 명소 갯수
# 오른쪽 서브트리
segt[13] = 0 5~5의 명소 갯수

# segt[12] 값과 찾고자하는 값이 같으므로 왼쪽 서브트리로 이동, 값은 유지

node = 12
start = 4
end = 4

# start == end 이므로 탐색 종료.

종료 후 최종값은 left = right = 4가 된다.
원래의 리스트값과 비교해보면 [1, 0, 0, 1, (1), 0, 1, 0] => 4번 인덱스까지의 명소 갯수는 3, 최종값과 일치하는것을 확인할 수 있다!

 

요약하자면 원하는 값을 가진 위치를 찾을때까지 범위를 절반씩 나눠가며, 세그먼트 트리를 파고드는것이다.
명령어 1 ~ 4의 수행 내용은 다음과 같다.
(참고로 나는 명소도 1-based로 맞춰줬다.)

 

명령어 == 1

- 명소인경우 해제
해당 세그먼트 트리에 도달할때까지 거쳐가는 모든 세그먼트 트리에 -1
- 명소가 아닌경우 명소로 추가
해당 세그먼트 트리에 도달할때가지 거쳐가는 모든 세그먼트 트리에 +1

 

명령어 == 2

현재 위치를 주어진 값만큼 이동시킴

 

명령어 == 3

- 명소가 아예 없다면(segt[1] == 0) -1 출력
아니라면, 명소의 갯수를 카운트할 변수 cnt를 1로 설정 후 현재위치까지의 명소의 갯수를 카운트한다.
👉 왜 1로 설정하냐면 만약 cnt = 3이 될 경우, 현재위치까지의 명소의 갯수는 2개라는뜻이고 우리가 찾아야 할 명소는 3번째의 명소이기 때문이다. (0으로 설정하고 따로 처리해줘도 됨)
카운팅 후 cnt값이 전체 명소의 값보다 크다면, 현재 위치 이후에는 명소가 없다는 뜻이므로 첫번째 명소를 타겟으로 설정한다.(cnt = 1)

그리고 cnt값과 일치하는 명소의 위치를 찾을때까지 세그먼트 트리를 파고든다.

값을 구했다면 명소의 위치가 현재 위치보다 크고 작음에 상관없이, (명소의 위치 - 현재 위치) % N을 한 결과값이 다음 명소까지의 거리값이 된다.
이건 파이썬의 특성덕이다.

 

전체 코드

# 메모리: 53528KB / 시간: 2248ms
from sys import stdin


input = stdin.readline

N, Q = map(int, input().split())
A = list(map(int, input().split()))
landmark = sum(A)  # 명소의 갯수

# 명소 업데이트(세그먼트 트리 업데이트) 함수
def update(target, node, left, right):
    # 현재 노드 값 업데이트
    seqt[node] += P  # P: target을 명소로 추가한다면 1, 명소리스트에서 삭제한다면 -1

    if left == right:
        return

    mid = (left + right) // 2

    if target <= mid:  # 찾는 명소가 mid보다 작거나 같다면 왼쪽 서브트리로
        update(target, 2*node, left, mid)
    else:  # mid보다 크다면 오른쪽 서브트리로
        update(target, 2*node + 1, mid+1, right)


# 현재 위치까지의 명소의 갯수를 카운트하는 함수 (현재 위치 포함)
def count_landmark(node, left, right):
    global cnt

    if left == right:
        return

    mid = (left + right) // 2

    if curr <= mid:  # 현재 위치가 mid보다 작거나 같다면 왼쪽 서브트리로
        count_landmark(2*node, left, mid)
    else:  # mid보다 크다면 왼쪽 서브트리값을 cnt에 추가, 오른쪽 서브트리로 이동
        cnt += seqt[2*node]
        count_landmark(2*node + 1, mid+1, right)


# 다음 명소의 위치를 찾는 함수
def find_landmark(node, left, right):
    global cnt, ret

    if left == right:
        ret = left
        return

    mid = (left + right) // 2

    if cnt <= seqt[2*node]:  # cnt값이 왼쪽 서브트리의 값보다 작거나 같다면, 왼쪽 서브트리로 이동
        find_landmark(2*node, left, mid)
    else:  # 왼쪽 서브트리보다 크다면, 왼쪽 서브트리의 값을 빼고 오른쪽 서브트리로 이동
        cnt -= seqt[2*node]
        find_landmark(2*node + 1, mid+1, right)


# 세그먼트 트리 생성
seqt = [0] * (4*N)

# 기존 명소들을 세그먼트 트리에 저장
for i in range(N):
    if A[i]:
        P = 1
        update(i+1, 1, 1, N)

# 현재 위치: 1-based로 관리함
curr = 1

# 쿼리문 실행
for _ in range(Q):
    cmd, *x = map(int, input().split())

    # 1: 명소 toggle
    if cmd == 1:
        target = x[0]

        # 이미 명소로 지정된 장소라면 제거
        if A[target-1]:
            A[target-1] = 0
            landmark -= 1
            P = -1  # 제거해줘야 하므로 -1로 설정
        else:
            A[target-1] = 1
            landmark += 1
            P = +1
        update(target, 1, 1, N)

    # 2: 인간 이동
    elif cmd == 2:
        dis = x[0]

        curr = ((curr - 1 + dis) % N) + 1  # 현재 위치에서 dis만큼 이동

    else:
        # 명소가 아예 없다면 -1 출력 후 넘어감
        if not seqt[1]:
            print(-1)
            continue

        cnt = 1  # 현재 위치 이후의 cnt번째 명소
        count_landmark(1, 1, N)

        # 만약 카운팅 한 결과값이 실제 명소의 갯수+1 일경우 -> curr 이후 명소 X, 첫번째 명소가 제일 가까운 셈
        if cnt == landmark+1:
            cnt = 1

        ret = 0
        find_landmark(1, 1, N)
        print((ret - curr) % N)

덕분에 세그먼트 트리를 알게됐다. 고마운 녀석~