2020/03/04

[zerojudge] d713: 中位數

import sys
from heapq import *
# https://docs.python.org/zh-tw/3/library/heapq.html
heapL = []
heapR = []
'''
heapL 放 <= 中位數的值,
因為我取左半部,都會取最大值,所以我存入負數。
heapR 放 >= 中位數的值,
因為我取右半部,都會取最小值,所以我存入原數。

heappop     : 從 heap 取出並回傳最小的元素
heappush    : only push
heappushpop : push and pop  
'''
heapify(heapL)
heapify(heapR)
# heapify() 將一個 list 轉成一個 heap
count = 0  
# 讀取記數,判斷 element 數量奇偶
# 奇數時,讓左邊多一個
# 偶數時,讓兩邊一樣多

for s in sys.stdin: # 開始 input
    n = int(s)
    count += 1
    if(count == 1):
        heappush(heapL, -n)
        print(n)
        continue
        # first element, 直接印出, 存入負值
        # 直接 continue 讀取下一個值
    largestL = -heapL[0]
    # heapL[0] 是 heapL 的最小值,
    # 加負號 = 左邊最大值 
    if(count % 2 == 0): 
        # element count 偶數時
        if(n >= largestL):
            heappush(heapR, n)  
            # 把 n 存入右邊,2 側數量相等
            smallestR = heapR[0]
            # 查詢 heapR[0] , 這是右邊的最小值
            r = (largestL + smallestR) // 2
        else:
            heappushpop(heapL, -n)
            # 把 n 存入左邊, 同時刪除 largestL
            heappush(heapR, largestL)
            # 把 largestL 存入右邊,2 側數量相等
            smallestR = largestL
            largestL = -heapL[0]
            # 新的左側最大值
            r = (largestL + smallestR) // 2
    else:    
        # element count 奇數時
        smallestR = heapR[0]
        if(n > smallestR):
            # 此時 n 要放在右邊
            # 右邊多一個要移到左邊
            heappush(heapL, -(smallestR))
            # smallestR 直接塞左邊
            heappushpop(heapR, n)
            # 把 n 放右邊並刪除原本最小值
            r = -(heapL[0])
            # 左邊比右邊多一個, 最大的數就是中位數
        else:
            heappush(heapL, -(n))
            # 直接把值塞左邊
            r = -(heapL[0])
            # 左邊比右邊多一個, 最大的數就是中位數
    print(r)

沒有留言:

張貼留言