[Python] 백준 13925번: 수열과 쿼리 13(다이아 문제 첫 성공)

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

13925번: 수열과 쿼리 13

 

이 글은 Segment Tree의 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)

 

가장 큰 특징은, update가 3개로 나뉘어있다는 것이다. 이것을 일반화하기 위해서 많은 시간을 투자하였다.

 

매우 중요한 부분은 2가지로 나뉜다. 먼저 그 첫번째이다.

propagete 함수에서의 연산의 단순화를 고민해보자. 다음은 특정 노드를 lazy값을 참조하여 변경시킬 때의 연산을 일반화한 것이다.(일반적인 문제에서 seg_tree[index] = seg_tree[index] + (end-start+1)*lazy[index]를 해주는데 이것의 표기를 단순화하면 S+R(range라는 뜻)*L이다.)

 - 합 연산: S + R*value = S*1 + R*value

 - 곱 연산: S*value = S*value + R*0

 - 대입 연산: R*value = S*0 + R*value

포인트는 1, value, 0, 그리고 value, 0, value에 대한 정보를 lazy1, lazy2로 나누어서 저장해주겠다는 것이다.

이것을 위해, lazy에 직접 값을 할당해주는 부분에서 3개로 분기해준다. 관련 코드는 update 함수의 아래 부분이다.

if left <= start and end <= right:
    if order == 1:
        lazy1[index] = 1
        lazy2[index] = value
    elif order == 2:
        lazy1[index] = value
        lazy2[index] = 0
    else:
        lazy1[index] = 0
        lazy2[index] = value

    propagate_segment_tree(index, start, end)

    return

 

앞서 언급하였던 중요하다는 것들 중에, 이제 그 두번째이다.

이제 lazy값이 중첩되지 않았을 때는 해결되었다. lazy 값이 중첩될 때가 큰 난관이었다.

 

노드 T에 lazy가 각각 A, B이 있으며, 이 노드의 부모에 대해 합 연산이 이루어져서 노드 T에 부모의 lazy값인 C, D를 전파해야 한다고 하자.

첫번째로, 부모 노드에 합 연산이 이루어졌을 때(C=1, D=value), 노드 T는 나중에, (A*S+B*R)+D*R이라는 값을 가져야 한다. 즉, A*S+(B+D)*R이 되어야하며, 이것을 통해 알 수 있는 것은 lazy2는 더하여 전파되어야 한다는 것이다. 그러니까, 합 연산에서 lazy에 D(lazy2 값)을 더하여 보내주면, 나중에 이것을 유효하게 사용할 수 있다는 것이다.

두번째로, 곱 연산이 이루어졌을 때(C=value, D=0)이다. 노드 T는 (A*S+B*R)*C = (A*C)*S+(B*C)*R라는 값을 가져야한다. 기존 값들에 모두 C가 곱해졌다. 즉, lazy1은 자식 lazy1, lazy2에 대해 곱하여 전해져야한다.

마지막으로, 대입 연산이 이루어졌을 때(C=0, D=value)는 0*S+D*R 값을 지니게 된다.

 

이런 예시로 알 수 있는 점은 결국 두 가지다. lazy1은 자식 lazy1, 2의 값에 곱을 해주어야한다는 것, 그리고 lazy2는 자식 lazy2에 단순히 더해주면 된다는 것이다. 이 결론은 propagate 함수의 아래 코드에서 나타난다.

if start != end:
    lazy1[index*2] = (lazy1[index*2]*lazy1[index]) % MOD
    lazy1[index*2+1] = (lazy1[index*2+1]*lazy1[index]) % MOD
    lazy2[index*2] = (lazy2[index*2]*lazy1[index] + lazy2[index]) % MOD
    lazy2[index*2+1] = (lazy2[index*2+1]*lazy1[index] + lazy2[index]) % MOD

 

이러한 발상을 기반으로 코드를 작성하면 된다. lazy1은 곱에 대한 리스트이므로 1이 default값이라는 것을 명심하자.

 

코드는 아래와 같다.

from sys import stdin
from math import ceil, log

input = stdin.readline
MOD = 10**9+7


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


def initialize_segment_tree(index, start, end):
    if start == end:
        seg_tree[index] = nums[start]
        return

    mid = (start + end)//2
    initialize_segment_tree(index*2, start, mid)
    initialize_segment_tree(index*2+1, mid+1, end)
    seg_tree[index] = (seg_tree[index*2] + seg_tree[index*2+1]) % MOD


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

    if right < start or end < left:
        return

    if left <= start and end <= right:
        if order == 1:
            lazy1[index] = 1
            lazy2[index] = value
        elif order == 2:
            lazy1[index] = value
            lazy2[index] = 0
        else:
            lazy1[index] = 0
            lazy2[index] = value

        propagate_segment_tree(index, start, end)

        return

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


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

    if right < start or end < left:
        return 0

    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)) % MOD


def propagate_segment_tree(index, start, end):
    if not (lazy1[index] == 1 and lazy2[index] == 0):
        seg_tree[index] = seg_tree[index]*lazy1[index] + \
            (end - start + 1)*lazy2[index]

        if start != end:
            lazy1[index*2] = (lazy1[index*2]*lazy1[index]) % MOD
            lazy1[index*2+1] = (lazy1[index*2+1]*lazy1[index]) % MOD
            lazy2[index*2] = (lazy2[index*2]*lazy1[index] + lazy2[index]) % MOD
            lazy2[index*2+1] = (lazy2[index*2+1]*lazy1[index] + lazy2[index]) % MOD

        lazy1[index] = 1
        lazy2[index] = 0


if __name__ == '__main__':
    N = int(input())
    nums = [-1] + list(map(int, input().split()))

    seg_length = get_segment_tree_length()
    seg_tree = [0]*seg_length
    lazy1 = [1]*seg_length
    lazy2 = [0]*seg_length
    initialize_segment_tree(1, 1, N)

    M = int(input())
    for _ in range(M):
        cur = list(map(int, input().split()))
        if cur[0] == 1 or cur[0] == 2 or cur[0] == 3:
            update_segment_tree(1, 1, N, cur[1], cur[2], cur[0], cur[3])
        else:
            print(query_segment_tree(1, 1, N, cur[1], cur[2]))

 

 

 

시간, 공간복잡도는 평범~느린 편에 속한다. 참고로 Python으로 성공한 사람은 없으며 Pypy3으로 제출해야한다.