For each value x in the array, poison occurs from time x to x+duration. We want the length of the union of these intervals.
For each value x in the array, lets naively add the duration D to the answer, and try to correct it. We have to subtract the overlap of this segment with the previous segment. The value 'prev' records the rightmost endpoint of the previous segment, so the overlap is prev - x. ('x' is the leftmost endpoint of the new segment.)
When these segments don't overlap (such as when x > prev), then we shouldn't remove anything. We handle this with max(0, ...) on the length of the segment.
def findPosisonedDuration(self, A, D): ans = 0 prev = float('-inf') for x in A: ans += D - max(0, prev - x) prev = x + D return ans
O(N log N) solution, but more straightforward:
Again, we want the length of the union of these intervals. We remember this is a classic problem very similar to Leetcode#56 - Merge Intervals.
For each interval, add the start point and end point as events and process them from left to right.
Now we keep 'bal' (balance), the count of how many start points we've seen minus the number of end points we've seen.
If the balance was positive, then between now and the previous point of interest, that segment was part of the union, and we should count it.
def findPosisonedDuration(self, A, duration): events =  OPEN, CLOSE = range(2) for x in A: events.append((x, OPEN)) events.append((x+duration, CLOSE)) events.sort() ans = bal = 0 prev = float('-inf') for x, typ in events: if bal > 0: ans += x - prev prev = x bal += 1 if typ == OPEN else -1 return ans