[백준/1572] 중앙값 (Python)

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

1572번: 중앙값

 

Difficulty : Platinum 5

 

Status : Solved

 

Time : 00:34:21

 

 

---

 

문제 설명

 

 

중앙값이란, 수열을 정렬했고, 그 크기가 N일 때, 1부터 시작해서 (N+1)/2번째 있는 원소가 그 수열의 중앙값이다. 예를 들어, {1, 2, 6, 5, 4, 3}에서는 3이고, {11, 13, 12, 15, 14}에서는 13이다.

오세준은 1초에 온도를 하나씩 재는 온도계를 만들었다. 이 온도계에는 작은 디스플레이 창이 하나 있는데, 이 창에는 지금부터 최근 K초 까지 온도의 중앙값을 표시해 준다. (온도를 재기시작한지 K초부터 표시한다. 그 전에는 아무것도 출력되지 않는다.)

오세준은 온도를 N초동안 쟀다. 그 시간 동안 온도계의 디스플레이 창에 뜨는 숫자의 합을 구하는 프로그램을 작성하시오.

다른 말로 하면, 길이가 N인 수열이 주어진다. 이 수열은 N-K+1 개의 길이가 K인 연속된 부분 수열이 존재한다. 이 부분 수열의 중앙값의 합을 출력하는 프로그램을 작성하시오.

 

 

입력 및 출력

 

입력

 

첫째 줄에 N과 K가 주어진다. N은 250,000보다 작거나 같은 자연수이고, K는 5,000보다 작거나 같은 자연수이다. N은 항상 K보다 크거나 같다. 둘째 줄부터 N개의 수가 한 줄에 하나씩 주어진다. 입력으로 주어지는 수는 65536보다 작거나 같은 자연수 또는 0이다.

 

출력

 

첫째 줄에 정답을 출력한다. 정답은 2^63-1보다 작거나 같다.

 

입력 예시

 

10 3
3
4
5
6
7
8
9
10
11
12

 

출력 예시

 

60

 

 

---

 

풀이

 

세그먼트 트리를 사용해 보자. 온도를 기준으로, 카운팅하는 세그먼트 트리를 고려해보면 된다.

 

• i번째 시점에 세그먼트 트리에는 [i-K:i]까지의 온도가 저장되어 있는 셈이다. 들어오는 값은 +1로, 나가야 할 값은 -1로 카운팅을 업데이트하면 된다.
• 이제 중앙값을 찾아야 할 텐데, 이는 이분 탐색을 바탕으로 진행 가능하다. 루트 노드를 기준으로 현재 탐색해야 할 순서값 idx를 잡자. 이 순서값이 왼쪽에 있는 카운팅 개수 left 내에 존재한다면, 왼쪽 서브트리를 탐색하면 된다. 그렇지 않다면 오른쪽 서브트리를 탐색하면 된다(이 때 탐색해야 할 순서값은 idx - left로 업데이트된다.)

 

 

풀이 코드

import sys
input = sys.stdin.readline
MAX = 65537

N, K = map(int, input().split())
temp_list = [int(input()) for _ in range(N)]

class SegTree :
  def __init__(self, ):
    self.tree = [0]*(MAX*4)

  def update(self, temp, val) :
    self._update(0, MAX, 1, temp, val)

  def _update(self, start, end, idx, target, val) :
    if target < start or target > end :
      return
    self.tree[idx] += val
    if start == end :
      return
    mid = (start + end) // 2
    self._update(start, mid, idx*2, target, val)
    self._update(mid+1, end, idx*2+1, target, val)

  def search(self) :
    return self._search(0, MAX, 1, (K + 1) // 2 )

  def _search(self, start, end, idx, val) :
    if start == end :
      return start

    mid = (start + end) // 2
    left = self.tree[idx*2]
    if left >= val :
      return self._search(start, mid, idx*2, val)
    else :
      return self._search(mid+1, end, idx*2+1, val-left)


segtree = SegTree()
ans = 0

for i in range(N) :
  segtree.update(temp_list[i], 1)
  if i > K-1 :
    segtree.update(temp_list[i-K], -1)
  if i >= K-1 :
    ans += segtree.search()

print(ans)

풀이 완료!