[백준/22976] 계산 최적화 (Python)

Problem : https://www.acmicpc.net/problem/22976

22976번: 계산 최적화

 

Difficulty : Platinum 3

 

Status : Solved

 

Time : 00:43:50

 

---

 

문제 설명

 

현욱은 덧셈과 곱셈을 계산할 수 있는 간단한 계산 장치를 만들었다. 이 계산 장치는 처음에 값 0에서 시작해서 N(1 <= N <= 5 * 10^5 )번의 연산을 순서대로 수행한 결과값을 반환한다. 연산은 아래의 두 종류가 있다.

 

• + k : 이전 값에 k를 더한다(1 <= k <= 10^9)
• * k : 이전 값에 k를 곱한다 (1 <= k <= 10^9)

연산 결과값이 아주 커질 수 있기 때문에 이 계산 장치는 항상 계산 결과값을 10^9+7로 나눈 나머지를 저장한다.
현욱은 이 계산 장치에서 일부 연산을 변경한 후 계산 결과가 어떻게 달라지는지를 확인하고 싶다. 현욱은 총 Q(1 <= Q <= 5 * 10^5 ) 번 계산 장치의 연산을 변경할 예정이며, 연산 변경은 다음과 같은 형태로 주어진다.

 

• i + k : i번째 연산을 이전 값에 k를 더하는 연산으로 변경한다 (1 <= i <= N, 1 <= k <= 10^9)
• i * k : i번째 연산을 이전 값에 k를 곱하는 연산으로 변경한다 (1 <= i <= N, 1 <= k <= 10^9)

모든 변경 사항은 누적된다.

현욱은 연산 변경이 하나 일어날 때마다 다시 처음부터 계산을 하는게 너무 비효율적이라고 생각한다. 현욱을 도와 현욱의 계산 장치를 최적화해서, 연산 변경이 일어날 때 결과를 빠르게 다시 계산하는 프로그램을 작성해보자. 

 

 

입력 및 출력

 

입력

 

첫줄에 연산의 개수 N과 연산을 변경할 횟수 Q가 공백으로 구분되어 주어진다(1 <= N, Q <= 5 * 10^5 ).

둘째 줄부터 N줄에 걸쳐 맨 처음 계산 장치를 구성하는 N개의 연산이 순서대로 주어진다.

그 다음부터 Q줄에 걸쳐 계산 장치의 연산을 어떻게 변경했는지 정보가 순서대로 주어진다.

 

출력

 

Q줄에 걸쳐 연산이 바뀔 때마다 계산 장치의 계산 결과값이 어떻게 바뀌는지 출력한다.

 

입력 예시

 

6 10
+ 3
* 2
+ 2
* 3
+ 4
* 5
1 + 2
3 * 1
4 * 5
2 + 6
2 * 2
1 * 2
3 * 1
4 + 4
1 + 5
3 * 4

 

출력 예시

 

110
80
120
220
120
20
20
40
90
240

 

 

---

 

풀이

 

이 문제에서 영감을 받았다.

 

2024.01.17 - [PS/백준] - [백준/13925] 수열과 쿼리 13 (Python)

[백준/13925] 수열과 쿼리 13 (Python)

 

이 문제에서 못다한 말을 하자면, 사칙연산은 일차식 함수의 꼴로 구성할 수 있고, 일차식 함수의 합성함수 역시 일차식이다.

 

f(x) = ax + b이고, 이를 [a, b]꼴로 저장한다고 가정하자.

• 리프 노드에서의 f(x)는 다음과 같다.
  • + k : [1, k]
  • * k : [k, 0]
  • 업데이트 시 리프 노드를 위 방식과 같이 업데이트할 수 있다.
• 또한, 리프 노드가 아닌 노드에서를 고려하자.
  • 왼쪽 노드의 함수를 f(x), 오른쪽 노드의 함수를 g(x) = cx + d라고 두자. 그렇다면 이 노드의 함수 h(x) = g(f(x))가 될 것이다.
  • 그렇다면 f(x) = ax + b이므로 h(x) = c(ax + b) + d = acx + bc + d 로 풀이할 수 있으며, 이 역시 일차식이다.
  • h(x)를 앞선 경우와 같이 저장하면 [ac, bc + d]가 된다.
• 즉, 우리는 연산 업데이트를 O(logN)의 시간복잡도로 수행할 수 있다. 그 결과는 세그먼트 트리의 루트 노드에 저장되며, 초깃값은 0이므로 우리가 찾는 값은 root[0] * 0 + root[1] = root[1]이 된다.

 

 

풀이 코드

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

N, Q = map(int, input().split())
ops = [input().split() for _ in range(N)]
index = [0]*N
tree = [[0]*2 for i in range(4*N)]

def init(start=0, end=N-1, idx=1) :
  if start == end :
    op, num = ops[start]
    if op == '+' :
      tree[idx] = [1, int(num)]
    else :
      tree[idx] = [int(num), 0]
    index[start] = idx
    return
  mid = (start + end) // 2
  init(start, mid, idx*2)
  init(mid+1, end, idx*2+1)
  tree[idx][0] = (tree[idx*2][0] * tree[idx*2+1][0]) % MOD
  tree[idx][1] = (tree[idx*2][1] * tree[idx*2+1][0] + tree[idx*2+1][1]) % MOD

def update(idx, op, num) :
  idx = index[idx]
  if op == '+' :
    tree[idx] = [1, num]
  else :
    tree[idx] = [num, 0]
  while idx > 1 :
    idx //= 2
    tree[idx][0] = (tree[idx*2][0] * tree[idx*2+1][0]) % MOD
    tree[idx][1] = (tree[idx*2][1] * tree[idx*2+1][0] + tree[idx*2+1][1]) % MOD
    

init()
for _ in range(Q) :
  idx, op, num = input().split()
  update(int(idx)-1, op, int(num))
  print(tree[1][1])

풀이 완료!