백준 13925 - 수열과 쿼리 13 [Python]

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

13925번: 수열과 쿼리 13

Lazy Propagation in Segment Tree

 

 

1번 쿼리는 (ax + b) * 1 + v = ax + b + v

2번 쿼리는 (ax + b) * v + 0= (v * a)x + v * b

3번 쿼리는 (ax + b) * 0 + v = v

꼴로 갱신할 수 있다.

 

곱하는 값, 더하는 값에 대해 2개의 lazy 배열이 필요하다.

 

 

템플릿 코드

import sys, math
import heapq
 
from collections import deque
input = sys.stdin.readline
 
hqp = heapq.heappop
hqs = heapq.heappush
 
#input
def ip(): return int(input())
def sp(): return str(input().rstrip())
 
def mip(): return map(int, input().split())
def msp(): return map(str, input().split().rstrip())
 
def lmip(): return list(map(int, input().split()))
def lmsp(): return list(map(str, input().split().rstrip()))
 
#gcd, lcm
def gcd(x, y):
    while y:
        x, y = y, x % y
    return x
 
def lcm(x, y):
    return x * y // gcd(x, y)
 
#prime
def isPrime(x):
    if x<=1: return False
    for i in range(2, int(x**0.5)+1):
        if x%i==0:
            return False
    return True
 
# Union Find
# p = {i:i for i in range(1, n+1)}
def find(x):
    if x == p[x]:
        return x
    q = find(p[x])
    p[x] = q
    return q
 
def union(x, y):
    x = find(x)
    y = find(y)
 
    if x != y:
        p[y] = x
 
def getPow(a, x):
    ret = 1
    while x:
        if x&1:
            ret = (ret * a) % MOD
        a = (a * a) % MOD
        x>>=1
    return ret
 
def getInv(x):
    return getPow(x, MOD-2)

 

정답코드. f는 곱셈, g는 덧셈이다.

h는 귀찮음 문제로 따로 빼뒀다.

MOD = int(1e9 + 7)
seg = [0 for _ in range(404040)]
t = [1 for _ in range(404040)]
d = [0 for _ in range(404040)]
 
def f(x, p):
    return x * p % MOD
 
def g(x, p):
    return (x + p) % MOD
 
def h(x):
    t[x] = f(t[x], t[x >> 1])
    d[x] = f(d[x], t[x >> 1])
    d[x] = g(d[x], d[x >> 1])
 
def init(x, s, e):
    if s == e:
        seg[x] = a[s - 1]
        return seg[x]
    m = s + e >> 1
    seg[x] = (init(x * 2, s, m) + init(x * 2 + 1, m + 1, e)) % MOD
    return seg[x]
 
def update(x, l, r, s, e, time, dif):
    updateLazy(x, s, e)
    if e < l or r < s:
        return
    if l <= s and e <= r:
        t[x] = f(t[x], time)
        d[x] = f(d[x], time)
        d[x] = g(d[x], dif)
        updateLazy(x, s, e)
        return
    m = s + e >> 1
    update(x * 2, l, r, s, m, time, dif)
    update(x * 2 + 1, l, r, m + 1, e, time, dif)
    seg[x] = (seg[x * 2] + seg[x * 2 + 1]) % MOD
 
def updateLazy(x, s, e):
    if s != e:
        h(x * 2)
        h(x * 2 + 1)
    seg[x] = (f(t[x], seg[x]) + (e - s + 1) * d[x]) % MOD
    t[x] = 1
    d[x] = 0
 
def getSum(x, l, r, s, e):
    updateLazy(x, s, e)
    if e < l or r < s:
        return 0
    if l <= s and e <= r:
        return seg[x] % MOD
    m = s + e >> 1
    return (getSum(x * 2, l, r, s, m) + getSum(x * 2 + 1, l, r, m + 1, e)) % MOD
 
n = ip()
a = lmip()
 
init(1, 1, n)
 
for i in range(ip()):
    q = lmip()
    if q[0] == 1:
        update(1, q[1], q[2], 1, n, 1, q[3])
    elif q[0] == 2:
        update(1, q[1], q[2], 1, n, q[3], 0)
    elif q[0] == 3:
        update(1, q[1], q[2], 1, n, 0, q[3])
    else:
        print(getSum(1, q[1], q[2], 1, n) % MOD)