[백준 14504번] 수열과 쿼리 18 (Python)

import sys
input = sys.stdin.readline

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

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

N = int(input()); M = 1<<18; sq = int(N**0.5)
seq = [0,*map(int,input().split())]
Q = int(input())
query = [[*map(int,input().split())] for i in range(Q)]
S = sorted({*seq}|{query[i][-1] for i in range(Q)},reverse=1)
Index = {S[i]:i+1 for i in range(len(S))}
fen = [[0]*M for i in range(N//sq+1)]
for i in range(N+1):
  update(i//sq,Index[seq[i]],1)
for q,*a,k in query:
  if q==1:
    L,R = a
    r = sum(SUM(n,Index[k]-1) for n in range(L//sq,R//sq))-sum(int(seq[i]>k) for i in range(L//sq*sq,L))+sum(int(seq[i]>k) for i in range(R//sq*sq,R+1))
    print(r)
  else:
    [i] = a
    update(i//sq,Index[seq[i]],-1)
    update(i//sq,Index[k],1)
    seq[i] = k

제곱근 분할 + 팬윅트리로 해결하였다.

제곱근 분할법에 대해서 처음으로 공부한 문제이다. 제곱근 분할에 대한 설명은 Sqrt Decomposition(제곱근 분할법) (tistory.com)을 참조하였다.

 

풀이를 설명하면,

1. 쿼리 전체를 미리 입력받는다. 그리고 원래 수열 A와 쿼리의 모든 k를 역순으로 좌표압축한다.

2. 구간을 제곱근 분할하고, 각 구간에 팬윅트리를 만든다. (fen[n] = n번째 구간의 팬윅트리)

3. 각 구간에 대해서 압축한 좌표를 update한다.

4. 1에서 저장한 쿼리를 순서대로 실행한다.

5. 1번쿼리의 경우 제곱근 구간 밖에 대해서는 k보다 큰 수를 직접 세어주고, 구간내에 대해서는 팬윅트리를 통해 k보다 큰 수를 세어준 뒤 합을 출력한다.

6. 2번쿼리의 경우 소속된 제곱근 구간의 팬윅트리를 업데이트하고, 수열도 업데이트한다.

 

시간복잡도는 각 1번쿼리마다 각 구간에 대해 팬윅트리를 최대 sqrt(N)번 실행하므로 시간복잡도는 O(NlogNsqrtN)이다.

 

제곱근 분할법을 공부할 수 있는 좋은 문제였다.