[
heap
array
]
BinarySearch 0772 Sum of Odd Length Medians
Problem statement
https://binarysearch.com/problems/Sum-of-Odd-Length-Medians/
Solution
Similar to Leetcode 0295. Find Median from Data Stream, but here we need to apply it for all A[i:]
.
Complexity
It is O(n^2 log n)
for time and O(n)
for space.
Code
class Solution:
def solve(self, A):
def RollingMedian(A):
small, large, ans = [], [], 0
for i, num in enumerate(A):
heappush(small, -num)
heappush(large, -heappop(small))
if len(small) < len(large):
heappush(small, -heappop(large))
if i % 2 == 0: ans += -small[0]
return ans
return sum(RollingMedian(A[i:]) for i in range(len(A)))
Remark
There is also O(n^2)
time complexity solution, where the idea is for each element find how many times it is median.