백준 23326 - 홍익 투어리스트

문제 설명

제한사항 및 입출력

구역의 개수가 5 * 10^6이고 쿼리의 개수가 10^6로 꽤나 많다.

풀이

총 세 가지의 연산이 있다. 1번 연산은 명소 토글, 2번 연산은 이동, 3번 연산은 가장 가까운 다음 명소찾기.

첫 번째로 신경써야 할 연산은 시계방향으로 몇 칸 움직여야 명소를 찾을 수 있는지이다. 이 구현은 첫 번째 명소가 나올 때까지 그냥 돌려서는 시간초과가 난다. 만약 구역이 희소하거나 없다면 하나하나 구역을 찾는 연산은 5 * 10^6의 시간을 갖는다. 시간 제한이 1초이므로 단순히 생각해봐도 100번만 반복해도 시간초과가 난다. 여기서 시간을 줄여야하는데 가장 적합한 방법은 이진탐색이다. 그냥 배열로 장소를 표현하면서 명소를 토글하는 방식으로 1번 연산을 수행한다면 1번 연산은 수행하기 편하지만 이분탐색이 불가능해진다. 따라서 장소의 인덱스만 따로 정렬한 상태로 이분탐색을 돌려야 한다. 그렇다면 장소의 인덱스를 적절한 위치에 삽입해야하는데 장소 삽입/삭제를 O(n)보다 빠르게 하면서 정렬이 되어있는 자료구조가 필요하다. 하지만 파이썬에는 그런 자료구조가 없다... 당장 생각난건 AVL이나 레드블랙트리같은 Balanced BST이기 때문에 구현하기 쉬운 Treap으로 구현하여 풀이했다.

내 코드

import sys
import random
input = sys.stdin.readline
mis = lambda: map(int, input().split())
rand = random.Random(0xC0FFEE)

class Treap:
    __slots__ = ("key", "prio", "l", "r")
    
    def __init__(self, key):
        self.key = key
        self.prio = rand.randrange(1<<30)
        self.l = None
        self.r = None


def rotate_right(p: Treap):
    q = p.l
    p.l = q.r
    q.r = p
    
    return q

def rotate_left(p: Treap):
    q = p.r
    p.r = q.l
    q.l = p
    
    return q

def insert(p: Treap, key):
    if not p:
        return Treap(key)
    
    if key < p.key:
        p.l = insert(p.l, key)
        
        if p.l.prio < p.prio:
            p = rotate_right(p)
    elif key > p.key:
        p.r = insert(p.r, key)
        
        if p.r.prio < p.prio:
            p = rotate_left(p)

    return p

def erase(p: Treap, key):
    if not p: 
        return None
    
    if key < p.key:
        p.l = erase(p.l, key)
    elif key > p.key:
        p.r = erase(p.r, key)
    else:
        if not p.l and not p.r:
            return None
        
        if not p.r or (p.l and p.l.prio < p.r.prio):
            p = rotate_right(p)
            p.r = erase(p.r, key)
        else:
            p = rotate_left(p)
            p.l = erase(p.l, key)
            
    return p

def find(p: Treap, key):
    while p:
        if key < p.key:
            p = p.l
        elif key > p.key:
            p = p.r
        else:
            return True
        
    return False

def lower_bound(p: Treap, key):
    ret = None
    
    while p:
        if p.key >= key:
            ret = p.key
            p = p.l
        else:
            p = p.r
            
    return ret

def min_key(p: Treap):    
    while p.l:
        p = p.l
        
    return p.key


n, q = mis()
A = [*mis()]

root = None
for i, v in enumerate(A, 1):
    if v == 1:
        root = insert(root, i)

x = 1
for _ in range(q):
    query = [*mis()]

    if query[0] == 1:
        i = query[1]
        
        if find(root, i):
            root = erase(root, i)
        else:
            root = insert(root, i)
    elif query[0] == 2:
        x = ((x - 1 + query[1]) % n) + 1
    else:
        if root is None:
            print(-1)
        else:
            y = lower_bound(root, x)
            
            if y is not None:
                print(y - x)
            else:
                m = min_key(root)
                print(n - x + m)

자료구조를 직접 구현하다보니 너무 길어졌는데... 이런 정렬된 Set이 필요한 문제는 파이썬이 아닌 TreeSet을 지원해주는 언어로 풀자...

 

• 쿼리 개수 = Q
• 각 쿼리 처리 비용:
  • 토글 -> treap의 insert, delete 연산 O(log N)
  • 질의 -> treap의 lower_bound(), min_key() 연산 O(log N)

시간 복잡도: O(Q log N)

 

고찰

다른 풀이는 생각나지 않는다... 다음 명소를 찾는 방법을 이진탐색이 아닌 다른 방법으로 할 수 있다면 최적화가 가능할 것 같다.