[백준 1280번] 나무 심기 (Python3)

import sys
input = sys.stdin.readline

def updatefen(i,v,fen):
  while i<=M:
    fen[i] += v
    i += i&-i

def sumfen(i,fen):
  SUM = 0
  while i:
    SUM += fen[i]
    i -= i&-i
  return SUM

def cost(i):
  lcnt = sumfen(i,cntfen); rcnt = cntfen[M]-lcnt
  lsum = sumfen(i,dfen); rsum = dfen[M]-lsum
  return lcnt*i-lsum+rsum-rcnt*i

N = int(input()); M = 1<<18; mod = 10**9+7
cntfen = [0]*(M+1); dfen = [0]*(M+1)

x = int(input())+1
updatefen(x,1,cntfen); updatefen(x,x,dfen)

result = 1
for i in range(N-1):
  x = int(input())+1
  result *= cost(x); result %= mod
  updatefen(x,1,cntfen); updatefen(x,x,dfen)
print(result)

팬윅트리로 해결하였다.

나무 심는 비용을 어떻게 설정하는지가 중요하다. 현재 나무의 위치에서 왼쪽에 있는 나무의 개수와 오른쪽에 있는 나무의 개수를 다르게 취급해주어야 하는데, 이를 활용한다.

 

풀이를 설명하면,

1. 팬윅트리를 두개 만든다. 하나는 구간 내의 나무심는 거리의 합을 나타내는 트리이고, 나머지 하나는 구간 내의 나무 개수의 합을 나타내는 트리이다.

2. x 위치에 나무를 심을때, 나무를 심는 비용은 (x 오른쪽에 있는 나무 거리의 합)-(x 오른쪽에 있는 나무 개수)*x+(x 왼쪽에 있는 나무 개수)-(x 왼쪽에 있는 나무 거리의 합) 이다.

3. 입력받는 숫자에 +1한 값을 x로 두고 (팬윅트리는 노드번호 1부터 시작) 나무 심는 비용의 곱을 구한다.

 

확실히 구간 합을 구해야하는 문제는 세그먼트 트리보다 팬윅트리가 성능면에서나 구현난이도 면에서나 훨씬 좋은 것 같다. 이러다가 세그먼트 트리 구현하는 법 까먹는건 아닌가 모르겠다.

 

 

오늘의 교훈) 팬윅트리는 유용하다.