[SWEA, Python] 3000 중간값 구하기

김서영·2022년 8월 10일

코딩테스트 스터디

목록 보기
10/11

문제

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV-fO0s6ARoDFAXT

해결

최대 Heap과 최소 Heap을 각각 두어 해결해야 하는 문제이다.
중간값을 기준으로 중간값보다 더 작을 경우에는 더 작은 수들만 모아둔 최대 heap에 넣는다.
중간값보다 더 클 경우에는 더 큰 수들만 모아둔 최소 heap에 넣는다.

입력받는 수에 의하면 항상 홀수 개의 주어진 수들로 연산을 하게 되는데, 이때 heap들의 크기가 일치하지 않는다면, 중간값을 업데이트 해줘야한다.

나는 푸는 과정에서 계속 1문제만 맞길래, 계속 삽질을 했는데 알고보니, heap을 전역변수로 선언해두고, 각 테스트 케이스 때마다 초기화를 안해줘서 생긴 문제였다. 그래서 초기화도 잊지 말자!

코드

import heapq
lower = []
higher = []

def putInMaxHeap(val):
    heapq.heappush(lower,-1*val)

def getFromMaxHeap():
    val = heapq.heappop(lower)
    return -1*val

T = int(input())
for test_case in range(1, T + 1):
    n, m = map(int,input().split())
    ans = 0
    lower.clear()
    higher.clear()
    for _ in range(n):
        a,b = map(int,input().split())
        if a < m:
            putInMaxHeap(a)
        else:
            heapq.heappush(higher,a)

        if b < m:
            putInMaxHeap(b)
        else:
            heapq.heappush(higher,b)

        if len(lower) != len(higher):
            if len(lower) > len(higher):
                heapq.heappush(higher,m)
                m = getFromMaxHeap()
            elif len(lower) < len(higher):
                putInMaxHeap(m)
                m = heapq.heappop(higher)

        ans = (ans+m)%20171109        
    print("#"+str(test_case)+" "+str(ans))
profile
하지만 저는 이겨냅니다. 김서영이죠?

0개의 댓글