[백준 1395] 스위치 (Python)

문제

1395번: 스위치

준규네 집에는 총 N개의 스위치가 있고 이를 편하게 1번부터 N번까지 차례대로 번호를 매겼다. 그리고 준규의 취미는 이 스위치들을 켜고 끄는 것이다.

준규가 하는 스위치를 갖고 노는 일은 크게 두 가지이다. 하나는 A번부터 B번 사이의 스위치 상태를 반전시키는 것이고 다른 하나는 C번부터 D번 사이의 스위치 중 켜져 있는 상태의 스위치의 개수를 세는 것이다.

하지만 준규가 싫증을 느껴 우리가 이 귀찮은 일을 떠맡게 되었고 프로그래밍을 통해 일을 처리하도록 결정하였다.

문제 해결 전략

기본 전략 설정

문제는 어떤 배열이 주어질 때 특정 구간의 질문에 대해 답을 요구하는 유형의 문제이다. N개의 스위치에 대해 껐다 켜는 작업이 빈번히 발생하는 와중에 특정 구간에서 켜진 스위치의 수에 대해 조사가 필요하다.

N개 스위치에 대해 켜져 있는 스위치의 수를 적절히 관리하여 구간 문제에 빠르게 답하기 위해 구간 트리를 활용하고자 하였다.

구간 트리에서 각각의 노드가 관리하는 정보는 해당 구간에서 켜져 있는 스위치의 수로 하였다.

문제의 초기 상태가 스위치가 모두 껴진 상태이니 구간 트리를 모두 0으로 초기화한다.

알고리즘 설계: A번부터 B번 스위치 반전시키기

스위치 반전 작업에 대한 구간 트리 갱신은 다음 순서로 진행된다.

구간 트리 노드와 업데이트할 구간의 교집합이 없다면 아무 작업도 하지 않는다.
업데이트할 구간 내에 현재 구간이 포함된다면 값을 갱신한다.
1. 현재 노드의 구간 길이가 i, 노드의 값(현재 켜져있는 스위치의 수)을 j라고 할 때 갱신된 노드의 값은 i - j이다.
2. 두 개의 자식 노드에 대해 재귀적으로 갱신한다.
3. 노드의 현재 값과 이전 값의 차이를 부모 노드로 전파한다.
이외의 경우 다음과 같이 작업한다.
1. 두 개의 자식 노드에 대해 재귀적으로 갱신한다.
2. 두 개의 자식 노드로부터 반환된 차이값을 각각 현재 노드 값에 더한다.
3. 노드의 현재 값과 이전 값의 차이를 부모 노드로 전파한다.

알고리즘 설계: C번부터 D번 스위치 중 켜져 있는 스위치의 개수 세기

구간 내 켜져 있는 스위치 수에 대한 질의는 다음 순서로 진행된다.

구간 트리 노드와 업데이트할 구간의 교집합이 없다면 0을 리턴한다.
질의 구간 내에 현재 노드 구간이 포함된다면 해당 노드 값을 리턴한다.
이외의 경우 두 자식 노드의 리턴값을 더해서 리턴한다.

특이사항의 반영

문제에서 주의해야 할 부분은 스위치를 끄거나 켤때 특정 범위의 모든 스위치를 반전시킨다는 것이다. 기본적으로 배열의 한 개 원소가 갱신될 때 구간 트리에서 갱신되어야 할 노드의 개수는 O(lgN)개이지만, 복수의 원소에 대해 업데이트해야 하는 경우 구간 트리에서 최대 O(NlgN)개 노드의 갱신이 필요하다.

범위에 대한 업데이트에 최적화하기 위해 지연 전파(Lazy Propagation)를 구현하였다.

지연 전파의 핵심은 구간 트리 노드에 대해 스위치 반전이 필요할 때 이를 바로 적용하지 않고 반전 요청이 있다는 사실만 기록해두는 것이다.

이후에 스위치를 반전시키거나 켜져있는 스위치 개수를 세기 위해 해당 노드의 값에 접근이 필요할 때 이를 반영 및 전파하게 된다.

정답 코드

import sys

# Evaluation이 필요한 flip 작업이 있는지 확인하고, Evalution과 Propagation을 진행한다.
def eval(index, rstart, rend):
    if lazy[index]:
        if rstart != rend:
            lazy[index * 2] = not lazy[index * 2]
            lazy[index * 2 + 1] = not lazy[index * 2 + 1]
        rlen = rend - rstart + 1
        rtree[index] = rlen - rtree[index]
        lazy[index] = False

def _flip(start, end, index, rstart, rend):
    if start > rend or end < rstart:
        return 0
    else:
        # 구간 트리 노드의 값 접근하기 전에 Evaluation한다.
        eval(index, rstart, rend)

        if start <= rstart and end >= rend:
            lazy[index] = True  # 현재 구간이 flip 구간에 포함될 경우 Lazy 플래그만 True로 변경한다.
            rlen = rend - rstart + 1
            result = rlen - rtree[index]
            return result - rtree[index]  # 플립했을 때 변하는 값의 차이를 부모 노드로 전파한다.

        if not (start <= rstart and end >= rend):
            rmid = (rstart + rend) // 2
            prev = rtree[index]
            # 두 개 자식노드에서 리턴한 차이값을 반영한다.
            rtree[index] += _flip(start, end, index * 2, rstart, rmid)
            rtree[index] += _flip(start, end, index * 2 + 1, rmid + 1, rend)

            # 바뀐 값의 차이를 부모 노드로 전파한다.
            return rtree[index] - prev

def flip(start, end):
    _flip(start, end, 1, 0, n - 1)

def _count(start, end, index, rstart, rend):
    if start > rend or end < rstart:
        return 0
    else:
        # 구간 트리 노드의 값 접근하기 전에 Evaluation한다.
        eval(index, rstart, rend)

        # 현재 구간이 flip 구간에 포함될 경우 현재 노드 값 리턴한다.
        if start <= rstart and end >= rend:
            return rtree[index]

        else:
            # 두 개 노드 값을 더하여 리턴한다.
            rmid = (rstart + rend) // 2
            return _count(start, end, index * 2, rstart, rmid) + _count(start, end, index * 2 + 1, rmid + 1, rend)

def count(start, end):
    return _count(start, end, 1, 0, n - 1)

n, m = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]
rtree = [0 for _ in range(4 * n)]
lazy = [False for _ in range(4 * n)]

for _ in range(m):
    o, s, t = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]
    s -= 1
    t -= 1
    if o == 0:
        flip(s, t)
    else:
        print(count(s, t))