[백준 17353번] 하늘에서 떨어지는 1,2, ..., R-L+1개의 별 (Python3)

import sys
input = sys.stdin.readline

def update(i,j,v,s,e,x):
  if i==s and j==e:
    seg[0][x] += v
    seg[1][x] += 1
    return
  mid = (s+e)//2
  if j<=mid:
    update(i,j,v,s,mid,x*2)
  elif i>mid:
    update(i,j,v,mid+1,e,x*2+1)
  else:
    update(i,mid,v,s,mid,x*2); update(mid+1,j,v+mid+1-i,mid+1,e,x*2+1)

def SUM(i,s,e,x):
  global S
  S += seg[0][x]+seg[1][x]*(i-s)
  if s==e:
    return
  mid = (s+e)//2
  if i<=mid:
    SUM(i,s,mid,x*2)
  else:
    SUM(i,mid+1,e,x*2+1)    

N = int(input())
data = [0,*map(int,input().split())]
seg = [[0]*4*N,[0]*4*N]
for _ in range(int(input())):
  q,*a = map(int,input().split())
  if q-1:
    S = data[a[0]]; SUM(*a,0,N,1)
    print(S)
  else:
    update(*a,1,0,N,1)

세그먼트 트리로 해결하였다.

수열과 쿼리 21 [백준 16975번] 수열과 쿼리 21 (Python3) (tistory.com) 문제의 업그레이드 버전이다.

 

수열과 쿼리 21 문제가 그랬듯이 이 문제도 lazy propagation으로 푼 사람들이 많지만, 그걸 굳이 이용하지 않아도 충분히 풀 수 있다. update는 구간에 대해서 이루어지지만 출력은 구간이 아닌 정점에 대해서 이루어지기 때문이다.

 

풀이를 설명하면,

1. 세그먼트 트리를 두개를 만든다. seg[0]는 구간의 시작지점에 더해진 값을, seg[1]은 구간에 더해진 횟수를 의미한다.

2. a~b 구간에 별이 떨어지면, 세그트리를 순회하면서 a~b 구간의 seg[0]에 v를 더하고, seg[1]에 1을 더한다. 이때 v는 1로 시작해서 구간이 분리될때마다 a와 분리된 시작지점의 차이만큼을 더해준 값이다. (즉, 1~6 구간에 1을 더해야하는데 1~4 구간과 5~6으로 나누어진다면, 1~4 구간에는 1을 더하고, 5~6 구간에는 5를 더해준다)

3. i 인덱스의 값을 출력할 때는, i가 속한 구간을 순회하면서, 각 구간의 seg[0]값과 구간의 시작지점과 i의 차이에 해당구간 seg[1]을 곱한 값의 총합을 출력한다.

 

 

오늘의 교훈) 구간에 더하고 구간합을 출력하는 문제는 lazy propagation이 필수적이지만, 구간에 더하고 정점값을 출력하는 문제는 굳이 lazy propagation을 이용할 필요가 없다.