Why not use min heaps instead of max heaps? Its easier to understand. Or maybe I’m missing something?
Also, you can stop looping when there are no more candidates for a next interval.
Of course heaps don’t have a major advantage over sorted lists except perhaps making the code a tiny bit simpler.
from heapq import heappush, heappop
class Interval:
def __init__(self, start, end):
self.start = start
self.end = end
def find_next_interval(intervals):
by_start_time = []
by_end_time = []
result = [-1]*len(intervals)
for idx, interval in enumerate(intervals):
heappush(by_start_time, (interval.start, idx))
heappush(by_end_time, (interval.end, idx))
while by_end_time and by_start_time:
end_time, idx = heappop(by_end_time)
while by_start_time and by_start_time[0][0] < end_time:
heappop(by_start_time)
if by_start_time:
result[idx] = by_start_time[0][1]
return result