BOJ_2696) 중앙값 구하기

너 오늘 코드 짰니?·2023년 1월 3일
1
post-custom-banner

https://www.acmicpc.net/problem/2696

말 그대로 중앙값을 구하면 되는 직관적인 문제.
그러나 입출력 형식이 다소 까다롭고 시간복잡도도 생각해야 하는 문제여서 기록에 남기기로 했다.

1차 시도

처음에는 heapsort를 사용한 정렬방식을 생각해 보았다.
heapsort를 하는 과정에서 pop이 일어나기 때문에 매 홀수 차시마다 heap을 새로 만들어 주어야 했다.
끄집어낸 수들을 temp에서 저장하고 있다가 홀수 차시마다 temp를 다시 heap으로 옮기고, 다시 끄집어내면서 중앙값 찾고를 반복했다.
결국 풀리긴 했는데 시간초과로 탈락.

  • 일단 입력 데이터에 비례해 n2n\over2 만큼 중앙값 구하는 연산이 반복되고
  • 중앙값을 구할 때 n2n\over2개의 데이터가 2회 왕복하므로

이 알고리즘의 시간복잡도는 n2n^2 이다.

+) 굳이 이렇게 복잡하게 heap에 넣었다 뺐다 하면서 구현할 필요없이 파이썬 내장함수 sort() 써가지고 하면 안되나? 생각이 들었는데 sort() 함수 시간복잡도 NlogNNlogN 인거 생각하면 어차피 시간초과 날거 같아서 시도 안했다.

import sys
import heapq

T = int(sys.stdin.readline())
dataset = [list() for i in range(T << 1)]
for t in range(T):
    size = int(sys.stdin.readline())
    dataset[t << 1] = [size]
    for i in range(int(size / 10) + 1):
        dataset[(t << 1) + 1] = dataset[(t << 1) + 1] + (list(map(int, sys.stdin.readline().split())))


for t in range(T):
    answer = list()
    temp = list()
    heap = list()
    size = dataset[t << 1][0]
    print((size >> 1) + 1)
    for i in range(size):
        while temp:
            heapq.heappush(heap, temp.pop())
        heapq.heappush(heap, dataset[(t << 1) + 1][i])
        heapq.heapify(heap)
        if (1 & i) == 0:    # 홀수
            for j in range(i >> 1):
                temp.append(heapq.heappop(heap))
            answer.append(heap[0])
    for idx in range(len(answer)):
        if (idx % 10) == 9 or idx == len(answer) - 1:
            print(answer[idx], end = "\n")
        else:
            print(answer[idx], end = " ")

2차 시도

한참 생각해도 답이 안나와서 문제 유형을 체크해봤는데 자료구조 문제였다. 대충 힙을 쓰면 pop을 하면서 중앙값에 빠르게 도달할 수 있겠다 정도의 아이디어는 있었는데, 힙을 하나만 쓰려고 고정관념이 박혀있던 것이 문제였다.

중앙값을 유지하면서, 큰 수들의 집합과 작은 수들의 집합 두 개를 힙으로 만들어 운용한다면 아무런 연산 없이 중앙값을 바로바로 출력해낼 수 있었다.

  • big, small 두 개의 힙과 middle(int 형 중앙값 변수)를 운용한다.
  • middle을 기점으로 새로운 데이터가 들어올 때 마다 대소 여부를 판단해서 big 또는 small 힙으로 넣는다.
  • 이 때 big과 small 의 크기를 일정하게 유지하는 것이 포인트다. (그래야 middle이 진짜 딱 중앙값이 되겠지.)

big과 small의 크기를 일정하게 유지하기 위해서 짝수번째 차시에는 무조건 small쪽으로 넣고, 홀수번째 차시에는 무조건 big 쪽으로 넣었다. (새로 들어온 데이터나 middle 둘 중 하나가 들어가겠지.)

그래서 홀수번째 차시에는 middle 값을 출력하기만 하면 되는데, 그 전에 small과 big에서 하나씩 heappop을 해준 다음 재정렬 해주는 과정을 거쳤다.
-> 각 차시마다 무조건 정해진 방향으로만 heappush를 하는 과정에서 정렬이 조금씩 틀어지는 경우가 생기므로 이를 맞춰주기 위한 작업이 필요했다.

import sys
import heapq

T = int(sys.stdin.readline())
dataset = [list() for i in range(T << 1)]
for t in range(T):
    size = int(sys.stdin.readline())
    dataset[t << 1] = [size]
    for i in range(int(size / 10) + 1):
        dataset[(t << 1) + 1] = dataset[(t << 1) + 1] + (list(map(int, sys.stdin.readline().split())))


for t in range(T):
    answer = list()
    big = list()
    small = list()
    middle, current = 0, 0
    size = dataset[t << 1][0]
    print((size >> 1) + 1)
    for i in range(size):
        current = dataset[(t << 1) + 1][i]
        if i == 0:
            middle = current
            answer.append(middle)
            continue
        if (1 & i) == 0:    # 홀수
            if current > middle:
                heapq.heappush(big, current)
            else:
                heapq.heappush(big, middle)
                middle = current
            temp = list()
            # 중간값 출력 전에 재정렬 하는 부분
            # 이부분을 좀 더럽게 짜서 좀 더 깔끔하게 고칠 필요가 있어보인다.
            s = -heapq.heappop(small)
            b = heapq.heappop(big)
            heapq.heappush(temp, s)
            heapq.heappush(temp, b)
            heapq.heappush(temp, middle)
            heapq.heappush(small, -heapq.heappop(temp))
            middle = heapq.heappop(temp)
            heapq.heappush(big, heapq.heappop(temp))
            answer.append(middle)
        else:    # 짝수
            if current <= middle:
                heapq.heappush(small,  -current)
            else:
                heapq.heappush(small, -middle)
                middle = current
    for idx in range(len(answer)):
        if (idx % 10) == 9 or idx == len(answer) - 1:
            print(answer[idx], end = "\n")
        else:
            print(answer[idx], end = " ")

이렇게 하면 중앙값을 구하는데 아무런 연산이 들어가지 않으므로, 시간복잡도를 n으로 줄일 수 있었다.

profile
안했으면 빨리 백준하나 풀고자.
post-custom-banner

0개의 댓글