김낭만

재밌거나 필요하다고 생각하는 걸 만듭니다.

백준(BOJ) 22976 계산 최적화 파이썬(Python)

문제 링크

세그먼트 트리 응용 문제.
단순히 곱하기랑 더하기로 연산을 나누고 세그먼트 트리를 만들면 결합법칙이 성립하지 않는다(즉 앞에서부터 계산해야 한다).
세그먼트 트리는 주어진 식을 앞에서부터 계산하지 않고 반씩 나누어서 Divide and Conquer한다.
따라서 세그먼트 트리로 만들 수 있게 결합법칙이 성립하는 꼴로 바꾸어 주어야 한다.
그러기 위해서는 “곱한 뒤에 더하는 연산”을 한 덩어리로 볼 수 있다.
곱셈과 덧셈을 이용한 어떤 식의 조합이든 곱한 뒤에 더하는 꼴로 표현할 수 있다.
이를테면, *3, +5, *2, +2는 *2 *3, +5 *2 +2 = *6, +12로 바꿀 수 있다.
이런 방식으로 만들기 위해, 세그먼트 트리의 각 노드를 곱할 수와 더할 수의 튜플로 만들어줄 수 있다.
단순히 곱하는 수만 있는 경우 더할 수는 덧셈의 항등원인 0이, 반대의 경우에 곱할 수는 곱셈의 항등원인 1이 될 것이다.
예를 들어 +3은 *1, +3으로 표현할 수 있다는 뜻이다.
그 이후엔 일반적인 세그먼트 트리를 활용하면 된다.
단, 루트 노드만 가져오면 되기 때문에 구간합을 구하는 함수는 불필요하고, 연산도 단순한 단일 연산이 아니라 곱한 뒤에 더하기인 것만 주의하면 된다.
추가로, 수의 범위가 매우 크기 때문에 중간 중간 모듈러 연산을 통해 범위를 줄여줄 필요가 있다.
구하고자 하는 것은 모든 연산 이후 최종 결과인데, 이는 루트 노드의 더할 수에 해당한다.
왜 곱할 수가 아니라 더할 수냐면, 시작하는 수가 0이기 때문에 루트 노드가 의미하는 것은 결국 0 * (곱할 수) + (더할 수)이고, 이는 결국 더할 수이기 때문이다.
처음엔 재귀 알고리즘이라 별 생각 없이 setrecursionlimit(10 ** 6)을 넣었는데, pypy에서 메모리 초과가 뜨더라.
생각해보면 입력 n의 최대 크기는 500000이고, 세그먼트 트리는 초기화와 업데이트 모두 로그 재귀이기 때문에, 재귀 깊이를 굳이 늘려 줄 필요가 없다.

#!python


from sys import stdin
input = stdin.readline

MOD = 10 ** 9 + 7

# 세그먼트 트리 초기화
# node는 노드 인덱스, l은 구간 시작 인덱스, r은 구간 끝 인덱스
def init(node, l, r):
    # 구간에 들어있는 원소가 하나(리프 노드)
    if l == r:
        segtree[node] = arr[l]
    # 그 외의 경우(중간 노드)
    else:
        m = (l + r) // 2
        # lc는 왼쪽 자식, rc는 오른쪽 자식
        lc, rc = init(node*2, l, m), init(node*2+1, m+1, r)
        # lc : * a + b, rc : * c + d 일 때,
        # node : * ac + bc + d
        segtree[node] = ((lc[0] * rc[1]) + rc[0]) % MOD, (lc[1] * rc[1]) % MOD
    return segtree[node]

# 세그먼트 트리 업데이트
# node는 노드, l과 r은 구간 시작 / 끝 인덱스, i는 바꿀 숫자의 인덱스, val은 바꿀 값, op는 연산자(+ 혹은 *)
def update(node, l, r, i, val, op):
    # 바꿀 숫자가 해당 노드의 구간 안에 없는 경우
    if i < l or i > r:
        pass
    # 정확히 바꿀 숫자를 가리키는 리프 노드인 경우
    elif l == r:
        # 더하기 연산이면 더할 값은 val, 곱할 값은 1(항등원)
        if op == '+':
            segtree[node] = val, 1
        # 곱하기 연산이면 더할 값은 0(항등원), 곱할 값은 val
        else:
            segtree[node] = 0, val
    # 그 외의 경우, 즉 구간 내에 바꿀 숫자를 포함하는 중간 노드인 경우
    else:
        m = (l + r) // 2
        # 왼쪽 자식과 오른쪽 자식
        lc, rc = update(node*2, l, m, i, val, op), update(node*2+1, m+1, r, i, val, op)
        # init()과 동일하게 값 업데이트
        segtree[node] = ((lc[0] * rc[1]) + rc[0]) % MOD, (lc[1] * rc[1]) % MOD
    return segtree[node]

n, q = map(int, input().split())
# arr[i][0]은 i번째 숫자에서 더해주는 값, arr[i][1]은 i번째 숫자에서 곱해주는 값
arr = []
segtree = [(0, 0) for _ in range(n * 4)]

for _ in range(n):
    op, i = input().split()
    i = int(i)
    
    # 더하기 연산이면 더할 값은 i, 곱할 값은 1(항등원)
    if op == '+':
        arr.append((i, 1))
    # 곱하기 연산이면 더할 값은 0(항등원), 곱할 값은 i
    else:
        arr.append((0, i))

# 루트 노드부터 재귀적으로 세그먼트 트리 초기화
init(1, 0, n -1)

for _ in range(q):
    i, op, k = input().split()
    # 루트 노드부터 재귀적으로 세그먼트 트리 업데이트하고 그 결과 출력
    print(update(1, 0, n - 1, int(i) - 1, int(k), op)[0] % MOD)

2024

맨 위로 이동 ↑

2023

세그먼트 트리

개요 선형적인 자료구조에서는 값에 접근하는 데에 \(O(1)\)이면 충분하지만, 대신 부분합을 구하는 데에는 \(O(N)\)이 필요하다. 그렇다면 이 자료구조를 이진 트리로 구성하면 어떨까? 값에 접근하는 데에 걸리는 시간이 \(O(\lg N)\)으로 늘어나지만 대신 부분합을 구하...

벨만-포드 알고리즘

개요 다익스트라 알고리즘과 함께 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 다익스트라 알고리즘보다 느리지만, 음수 가중치 간선이 있어도 작동하며, 음수 가중치 사...

다익스트라 알고리즘

개요 다익스트라 알고리즘은 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘 중 하나이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 단, 다익스트라 알고리즘은 음수 가중치 엣지를 허용하지 않는다. 이 경우에는 벨만-...

맨 위로 이동 ↑