<p>The <strong>median</strong> is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle values.</p>

<ul>
	<li>For examples, if <code>arr = [2,<u>3</u>,4]</code>, the median is <code>3</code>.</li>
	<li>For examples, if <code>arr = [1,<u>2,3</u>,4]</code>, the median is <code>(2 + 3) / 2 = 2.5</code>.</li>
</ul>

<p>You are given an integer array <code>nums</code> and an integer <code>k</code>. There is a sliding window of size <code>k</code> which is moving from the very left of the array to the very right. You can only see the <code>k</code> numbers in the window. Each time the sliding window moves right by one position.</p>

<p>Return <em>the median array for each window in the original array</em>. Answers within <code>10<sup>-5</sup></code> of the actual value will be accepted.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,3,-1,-3,5,3,6,7], k = 3
<strong>Output:</strong> [1.00000,-1.00000,-1.00000,3.00000,5.00000,6.00000]
<strong>Explanation:</strong> 
Window position                Median
---------------                -----
[<strong>1  3  -1</strong>] -3  5  3  6  7        1
 1 [<strong>3  -1  -3</strong>] 5  3  6  7       -1
 1  3 [<strong>-1  -3  5</strong>] 3  6  7       -1
 1  3  -1 [<strong>-3  5  3</strong>] 6  7        3
 1  3  -1  -3 [<strong>5  3  6</strong>] 7        5
 1  3  -1  -3  5 [<strong>3  6  7</strong>]       6
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,2,3,4,2,3,1,4,2], k = 3
<strong>Output:</strong> [2.00000,3.00000,3.00000,3.00000,2.00000,3.00000,2.00000]
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= k &lt;= nums.length &lt;= 10<sup>5</sup></code></li>
	<li><code>-2<sup>31</sup> &lt;= nums[i] &lt;= 2<sup>31</sup> - 1</code></li>
</ul>

import heapq
from collections import defaultdict


class Solution:
    def find_median(self, max_heap, min_heap, heap_size):
        if heap_size % 2 == 1:
            return -max_heap[0]
        else:
            return (-max_heap[0] + min_heap[0]) / 2

    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        max_heap = []
        min_heap = []
        heap_dict = defaultdict(int)
        res = []
        
        # min_heap to store the larger half of the numbers, we want lowest number in this half each time
        # max_heap to store the smaller half of the numbers, store negated values to get max heap, we want highest number in this half each time
        for i in range(k):
            heapq.heappush(max_heap, -nums[i])
            heapq.heappush(min_heap, -heapq.heappop(max_heap))
            if len(min_heap) > len(max_heap):
                heapq.heappush(max_heap, -heapq.heappop(min_heap))
        
        median = self.find_median(max_heap, min_heap, k)
        res.append(median)
        
        for i in range(k, len(nums)):
            prev_num = nums[i - k]
            heap_dict[prev_num] += 1 # lazy removal: "1" is to mark the number as removed from heap
            
            # if prev_num <= median, it means previous number was in max_heap, set balance to -1
            # "-1" balance means after removing prev_num, max_heap will have one less element than min_heap
            if prev_num <= median:
                balance = -1
            # else, it means previous number was in min_heap, set balance to 1
            # "1" balance means after removing prev_num, max_heap will have one more element than min_heap
            else:
                balance = 1
            
            # if current number is less than or equal to median, it should be added to max_heap, increase balance by 1
            if nums[i] <= median:
                balance += 1
                heapq.heappush(max_heap, -nums[i])
            # else current number should be added to min_heap, decrease balance by 1
            else:
                balance -= 1
                heapq.heappush(min_heap, nums[i])
            
            # at this point, balance can be -2, 0 or 2
            # "0" means it's balanced, no need to do anything
            # "-2" means max_heap has two less elements than min_heap, need to move one element from min_heap to max_heap
            # "2" means max_heap has two more elements than min_heap, need to move one element from max_heap to min_heap
            if balance < 0:
                heapq.heappush(max_heap, -heapq.heappop(min_heap))
            elif balance > 0:
                heapq.heappush(min_heap, -heapq.heappop(max_heap))


            # after marking the number as removed, when the number is at the top of the heap, we need to actually remove it and mark it back as "0"
            while max_heap and heap_dict[-max_heap[0]] > 0:
                heap_dict[-max_heap[0]] -= 1
                heapq.heappop(max_heap)
            
            while min_heap and heap_dict[min_heap[0]] > 0:
                heap_dict[min_heap[0]] -= 1
                heapq.heappop(min_heap)

            median = self.find_median(max_heap, min_heap, k)
            res.append(median)
        
        return res

# Time complexity: O(n * log(k)), fist for outer loop O(k) with each heap operation O(log(k)) so O(k * log(k)), second for outer loop O(n - k) with each heap operation O(log(k)) so O((n - k) * log(k)), total O(k * log(k) + (n - k) * log(k)) = O(n * log(k))
# Space complexity: O(n), heap_dict worst case O(n)
29496