[Python] 백준 1395번: 스위치

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

1395번: 스위치

 

Segment Tree의 응용인 Lazy propagation의 응용 문제이다. Lazy propagation의 상세한 설명은 아래 링크에 담겨있다.

https://rccode.tistory.com/entry/Python-%EB%B0%B1%EC%A4%80-10999%EB%B2%88-%EA%B5%AC%EA%B0%84-%ED%95%A9-%EA%B5%AC%ED%95%98%EA%B8%B0-2

[Python] Lazy Propagation (백준 10999번: 구간 합 구하기 2)

 

특정 범위에 특정 값을 더하는 것이 아니라, 스위치를 On/Off 하는 방식이므로 이에 대해서 변형을 해줘야한다. 가장 중요한 propagate 함수를 보자.

 

def propagate_segment_tree(index, start, end):
    if lazy[index] != 0:
        seg_tree[index] = (end - start + 1) - seg_tree[index]

        if start != end:
            lazy[index*2] ^= 1
            lazy[index*2+1] ^= 1

        lazy[index] = 0

 

 

lazy의 값은 0 또는 1이며, 1일 때 스위치를 반전시키는 식으로 설계한다. (end - start + 1)은 범위에서 모든 스위치를 켰을 때의 합이고, 여기서 현재 값을 빼주면 결과적으로 스위치를 반전시켰을 때의 값을 얻게 된다.

그리고 자식 노드의 lazy 값을 반전시킨다. 1을 넣어주는 것이 아니라 반전시켜주는 이유는 아래 케이스를 통해 알 수 있다.

3 3

0 1 3
0 1 3

1 1 2

 

그리고, update 함수에서도 위와 같이 해주면 된다.

 

아래는 소스 코드이다.

 

import sys
import math
input = sys.stdin.readline


def get_segment_tree_length():
    if N & (N-1) == 0:
        return 2*N
    else:
        return pow(2, math.ceil(math.log(N, 2)) + 1)


def update_segment_tree(index, start, end, left, right):
    propagate_segment_tree(index, start, end)

    # Out of range
    if right < start or end < left:
        return

    # In range
    if left <= start and end <= right:
        seg_tree[index] = (end - start + 1) - seg_tree[index]

        if start != end:
            lazy[index*2] ^= 1
            lazy[index*2+1] ^= 1

        return

    mid = (start + end)//2
    update_segment_tree(index*2, start, mid, left, right)
    update_segment_tree(index*2+1, mid+1, end, left, right)
    seg_tree[index] = seg_tree[index*2] + seg_tree[index*2+1]


def query_segment_tree(index, start, end, left, right):
    propagate_segment_tree(index, start, end)

    # Out of range
    if right < start or end < left:
        return 0

    # In range
    if left <= start and end <= right:
        return seg_tree[index]

    mid = (start + end)//2
    return query_segment_tree(index*2, start, mid, left, right) + query_segment_tree(index*2+1, mid+1, end, left, right)


def propagate_segment_tree(index, start, end):
    if lazy[index] != 0:
        seg_tree[index] = (end - start + 1) - seg_tree[index]

        if start != end:
            lazy[index*2] ^= 1
            lazy[index*2+1] ^= 1

        lazy[index] = 0


if __name__ == '__main__':
    N, M = map(int, input().split())

    length = get_segment_tree_length()
    seg_tree = [0]*length
    lazy = [0]*length

    for _ in range(M):
        o, s, t = map(int, input().split())
        if o == 0:
            update_segment_tree(1, 1, N, s, t)
        else:
            print(query_segment_tree(1, 1, N, s, t))