[백준/16404] 주식회사 승범이네 (Python)

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

16404번: 주식회사 승범이네

 

Difficulty : Platinum 3

 

Status : Solved

 

Time : 00:22:01

 

---

 

문제 설명

 

승범이는 평소 래퍼 도끼를 흠모해왔지만, 도끼만큼 랩을 잘할 수 없다는 것을 깨닫고 도끼만큼 돈이라도 벌자는 결심을 한다. 그래서 휴학 후 ㈜승범이네를 창업했다.

㈜승범이네는 판매원들로만 이루어진 다단계 회사이다.

승범이를 제외한 모든 판매원은 사수가 배정되는데, 사수는 한 회원당 단 한 명씩만 배정된다. 만약 판매원 A가 B의 사수라면, B를 A의 부사수라고 부른다.

㈜승범이네의 수익구조는 기형적이다.

판매원들은 제품을 자비로 사서 판매한다. 이때 제품을 구매가격보다 저렴하게 판매하게 되면 손해를 보게 되는데, 어떤 회원 A가 손해를 보면 그 회원의 모든 부사수도 같은 만큼의 손해를 보게 된다. 그러면 부사수들의 부사수들도 손해를 보게 되고, 그들의 부사수들도 손해를 보게 되고, … ,결국 A와 A 밑의 모든 판매원이 A가 잃은 만큼의 손해를 보게 된다.

반대로 판매원 A가 제품을 비싸게 팔아 이익이 생길 경우, A와 A밑의 모든 판매원이 같은 이익을 얻을 수 있다.

승범이는 직원들이 현재 얼마만큼 돈을 벌었는지 감시하기 위해 다음 두 종류의 명령을 처리하려고 한다.

 

• 1 i w : 직원 i가 w만큼 이익/손해를 본다. (이익은 양수로, 손해는 음수로 주어진다)
• 2 i : 직원 i의 현재 통장 잔액을 출력한다.

 

직원들은 빈 통장을 갖고 일을 시작하며, 이익과 손해가 실시간으로 통장에 기록된다. 물론 통장 잔액은 음수일 수도 있다.

일을 시작하기 직전에 플래티넘 승급전을 남겨두고 온 것을 떠올린 승범이는 우리에게 일을 맡기고 집으로 달려가버렸다.

만년 골드 승범이를 위해 문제를 대신 해결해주자.

 

입력 및 출력

 

입력

 

첫 번째 줄에 승범이를 포함한 판매원들의 수 N(1 ≤ N ≤ 100,000), 명령의 수 M(1 ≤ M ≤ 100,000)이 주어진다. 판매원들은 1번부터 N번까지 번호가 매겨지며, 승범이는 항상 1번이다.

두 번째 줄에 판매원 1번부터 N번까지의 사수가 순서대로 공백으로 구분되어 주어진다. 승범이는 사수가 없으므로 -1이 주어진다.

세 번째 줄부터 M개의 줄에 걸쳐 위에서 설명한 명령(i, w는 정수, 1 ≤ i ≤ N, -10,000 ≤ w ≤ 10,000) 이 주어진다.

 

출력

 

2번 명령이 주어질 때마다 한 줄에 하나씩 해당하는 직원 i의 잔고 상황을 출력한다.

 

입력 예시

 

5 7
-1 4 1 1 1
1 3 10
2 3
1 1 -7
2 2
1 4 10
2 2
2 1

 

출력 예시

 

10
-7
3
-7

 

 

---

 

풀이

 

1번 쿼리를 주목하자. 우리는 트리 내의 i번 노드와 그 모든 후손들에 대해 전부 같은 값을 업데이트할 필요가 있다. 즉 이 1번 쿼리를 효율적으로(O(logN)정도) 처리하기 위해 Lazy Propagation을 가장 먼저 떠올릴 수 있다.

 

여기서 문제는, 이 Lazy를 처리하기에는 트리의 구조가 매우 어렵다는 점이다. 트리 특성상 구조의 기형성이 발목을 잡는다. 이를테면 최악의 경우(일자형 트리)는 Lazy 구조를 사용하더라도 O(N)의 시간복잡도까지 나올 수 있다. 이런 상황을 막기 위해, 트리를 배열 형태로 바꾸어 보면 어떨까? 라고 떠올렸다.

 

우선 배열의 구조는 간단하게 정의할 수 있다.

• 현재 시점의 노드들의 모든 자손 노드들은 그 노드의 오른쪽에 위치한다(역은 성립하진 않는다)
• 그렇다면 현재 시점의 노드를 a, 마지막에 방문한 자손 노드를 b라 두자. [a, b]구간의 모든 노드들은 a노드 그 자신이거나 그 노드의 후손이 된다.
• i번째 사원이 a번 노드를 가리킨다고 하자(employee[i] = a) 노드를 업데이트하는 것은, [a, b] 구간 전체에 대한 업데이트 문제로 바뀐다. 즉 Lazy Segment Tree 문제로 치환되었다! 

...이를 처리하기 위해, 우리는 모든 트리를 순차적으로 방문하며 방문 순서를 매길 필요가 있다. DFS를 수행하여 방문 지점을 기록하고, 마지막에 방문한 자손 노드의 순서를 저장하면 된다[a, b]. 그리고 이렇게 변형된 Lazy Segment Tree 문제를 룰루랄라 풀면 된다.

 

간단한 트릭이라 쉽게 생각해낼 수 있었는데, 이를 따로 오일러 경로 테크닉(ETT)이라고 따로 칭한다는 걸 나중에야 알았다.

 

 

풀이 코드

from collections import defaultdict
import sys
sys.setrecursionlimit(10000000)
input = sys.stdin.readline

N, M = map(int, input().split())
parent = list(map(int, input().split()))
children = defaultdict(list)

for i in range(1, N) :
  children[parent[i]-1].append(i)

employee_info = [[0]*2 for _ in range(N)]
eidx = -1
def dfs(node) :
  global eidx
  eidx += 1
  employee_info[node][0] = eidx
  for child in children[node] :
    dfs(child)
  employee_info[node][1] = eidx

class LazySegTree :
  def __init__(self) :
    self.tree = [0]*(4*N)
    self.lazy = [0]*(4*N)

  def propagate(self, start, end, idx) :
    if start != end :
      self.lazy[idx*2] += self.lazy[idx]
      self.lazy[idx*2+1] += self.lazy[idx]
    else :
      self.tree[idx] += self.lazy[idx]
    self.lazy[idx] = 0

  def update(self, left, right, val) :
    def _update(start, end, idx) :
      self.propagate(start, end, idx)
      if right < start or left > end :
        return
      if left <= start <= end <= right :
        self.lazy[idx] += val
        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, target) :
    start, end, idx = 0, N-1, 1
    while start < end :
      self.propagate(start, end, idx)
      mid = (start + end) // 2
      if target <= mid :
        end = mid
        idx *= 2
      else :
        start = mid+1
        idx = idx*2+1
    self.propagate(start, end, idx)
    print(self.tree[idx])

dfs(0)
tree = LazySegTree()
for _ in range(M) :
  q, *cmd = map(int, input().split())
  if q == 1 :
    i, val = cmd
    l, r = employee_info[i-1]
    tree.update(l, r, val)
  else :
    i = cmd[0]
    i = employee_info[i-1][0]
    tree.search(i)

풀이 완료!