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

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

17353번: 하늘에서 떨어지는 1, 2, ..., R-L+1개의 별

 

Difficulty : Platinum 2

 

Status : Solved

 

Time : 01:25:24

 

---

 

문제 설명

 

 

욱제의 은밀한 취미 중 하나는 매일 밤하늘을 감상하는 것이다. 😓 욱제는 하늘의 별들이 다음과 같은 규칙들을 따르며 떨어지는 걸 관찰했다.

1. 별이 떨어지는 위치는 N개의 점이다. 점은 순서대로 1, 2, ..., N의 번호를 갖는다.
2. 매일 밤 별들은 1, 2, ..., N의 연속한 부분 구간 [L, R]에 떨어진다.
3. [L, R]에 별이 떨어지면, 각 점에는 순서대로 1, 2, ..., R-L+1개의 별이 떨어진다. 다시 말해, L에는 1개, L+1에는 2개, ..., R에는 R-L+1개의 별이 떨어진다.

욱제는 하늘에서 떨어지는 별들을 기록하다가 잠이 들어버렸다!! 혹시나 했지만 역시나, 여러분은 욱제를 대신해 아래의 쿼리를 수행해야 한다. (ㅎㅎ;; ㅈㅅ.. ㅋㅋ!!)

 

• 1 L R: [L, R]에 별이 떨어진다. (1 ≤ L ≤ R ≤ N)
• 2 X: 점 X에 떨어진 별의 개수의 합을 출력한다. (1 ≤ X ≤ N)

 

 

 

입력 및 출력

 

입력

첫째 줄에 별이 떨어지는 점의 개수 N이 주어진다. (1 ≤ N ≤ 10^5)

둘째 줄에 욱제가 잠들기 전까지 세어 놓은, 이미 떨어진 별들의 개수 A1, ..., AN이 공백을 사이에 두고 주어진다. (0 ≤ A1, ..., AN ≤ 10^6)

셋째 줄에는 쿼리의 개수 Q가 주어진다. (1 ≤ Q ≤ 10^5)

넷째 줄부터 Q개의 줄에는 쿼리가 한 줄에 하나씩 주어진다.

 

출력

 

2번 쿼리에 대한 답을 한 줄에 하나씩 출력한다.

 

입력 예시

 

5
1 2 1 2 1
4
1 1 5
2 5
1 2 5
2 5

 

출력 예시

 

6
10

 

 

---

 

풀이

 

 

여러 모로 의미 깊은 문제였다...!

 

• 최초의 Platinum 2 문제이기도 하고
• 내가 3년 전에는 풀지 못했던 문제이기도 하다

3년 전의 나보다 성장했음을 느낀다 :)

 

각설하고, 이 문제의 구간 업데이트, 구간 합은 느리게 갱신되는 세그먼트 트리를 생각나게 만든다. 모든 쿼리가 O(logN)의 시간복잡도를 보여야 통과 가능하기 때문이다. 문제는 구간 업데이트인데, 구간 업데이트가 [L, R]에서 [1, R-L+1]로 가변적으로 진행되어야 한다는 점이다. 여기서 우리는 lazy 리스트를 주목해야 한다.

 

이 문제에서 lazy[idx] = [a, d]을 저장할 것이다. 이 값이 의미하는 것은

 

• 초항이 a이고
• 공차가 d인 등차수열

을 의미한다. 즉 일반해는 a + nd (n >= 0)이 될 것이다. 잘 생각해보면, 업데이트는 [1, 1]값으로 이루어짐을 알 수 있다.

 

또한, 등차수열과 등차수열의 각 항의 합 역시 등차수열을 이룸이 자명하다. (a + nd) + (b + nd') = (a + b) + n(d + d') 꼴로 정리할 수 있기 때문이다. 즉, [a, d] 정보가 위에서 전파된다면 덧셈의 형태로 현제 인덱스의 lazy 정보를 갱신할 수 있다.

 

등차수열의 연속 부분수열 역시 등차수열이다. lazy[idx] 가 [start, end]구간의 정보를 저장한다고 가정하자. (start < end) 그렇다면, 이를 전파할 때 lazy[idx*2], lazy[idx*2+1] 구간에 정보를 전달해야 하며, 그 구간은 mid = (start + end) // 2일때 각각 [start, mid], [mid+1, end]이다.

 

여기서 start 인덱스의 값이 a + (start - start )d = a라면, mid + 1의 인덱스의 값은 a  + (mid+1 - start)d가 성립한다. 즉

 

• lazy[idx*2]에는 [ a, d ]
• lazy[idx*2+1]에는 [ a  + (mid+1 - start)d , d ]

값이 전달됨을 알 수 있다! 이를 고려하여 코드를 작성하면 된다.

 

 

 

풀이 코드

import sys
input = sys.stdin.readline

N = int(input())
stars = list(map(int, input().split()))
Q = int(input())

class LazySegTree() :
  def __init__(self) :
    self.tree = [0]*(4*N)
    self.lazy = [[0, 0] for _ in range(4*N)]
    def _init(start, end, idx) :
      if start == end :
        self.tree[idx] = stars[start]
        return
      mid = (start + end) // 2
      _init(start, mid, idx*2)
      _init(mid+1, end, idx*2+1)
    _init(0, N-1, 1)

  def propagate(self, start, end, idx) :
    if not self.lazy[idx][1] :
      return
    if start < end :
      mid = (start + end) // 2
      lmid = (mid + 1 - start) * self.lazy[idx][1] + self.lazy[idx][0]
      self.lazy[idx*2][0] += self.lazy[idx][0]
      self.lazy[idx*2][1] += self.lazy[idx][1]
      self.lazy[idx*2+1][0] += lmid
      self.lazy[idx*2+1][1] += self.lazy[idx][1]
    else :
      self.tree[idx] += self.lazy[idx][0]
    self.lazy[idx] = [0, 0]
   

  def update(self, L, R) :
    def _update(start, end, idx) :
      self.propagate(start, end, idx)
      if end < L or R < start :
        return
      if L <= start <= end <= R :
        self.lazy[idx][0] += start - L + 1
        self.lazy[idx][1] += 1
        self.propagate(start, end, idx)
        return
      mid = (start + end) // 2
      _update(start, mid, idx*2)
      _update(mid+1, end, idx*2+1)
    _update(0, N-1, 1)

  def search(self, x) :
    def _search(start, end, idx) :
      self.propagate(start, end, idx)
      if start == end == x :
        print(self.tree[idx])
        return
      if x < start or x > end :
        return
      mid = (start + end) // 2
      if x <= mid :
        _search(start, mid, idx*2)
      else :
        _search(mid+1, end, idx*2+1)
    _search(0, N-1, 1)

segtree = LazySegTree()
for _ in range(Q) :
  q, *cmd = map(int, input().split())
  if q == 1 :
    L, R = cmd
    segtree.update(L-1, R-1)
  else :
    x = cmd[0]
    segtree.search(x-1)

풀이 완료!