|
| 1 | +from heapq import heappush, heappop |
| 2 | +from typing import List |
| 3 | + |
| 4 | + |
| 5 | +class Solution: |
| 6 | + def leftmostBuildingQueries(self, heights: List[int], queries: List[List[int]]) -> List[int]: |
| 7 | + max_idx = [] |
| 8 | + results = [-1] * len(queries) |
| 9 | + store_queries = [[] for _ in heights] |
| 10 | + |
| 11 | + for idx, (a, b) in enumerate(queries): |
| 12 | + if a < b and heights[a] < heights[b]: |
| 13 | + results[idx] = b |
| 14 | + elif a > b and heights[a] > heights[b]: |
| 15 | + results[idx] = a |
| 16 | + elif a == b: |
| 17 | + results[idx] = a |
| 18 | + else: |
| 19 | + store_queries[max(a, b)].append( |
| 20 | + (max(heights[a], heights[b]), idx) |
| 21 | + ) |
| 22 | + |
| 23 | + for idx, height in enumerate(heights): |
| 24 | + while max_idx and max_idx[0][0] < height: |
| 25 | + _, q_idx = heappop(max_idx) |
| 26 | + results[q_idx] = idx |
| 27 | + for element in store_queries[idx]: |
| 28 | + heappush(max_idx, element) |
| 29 | + |
| 30 | + return results |
| 31 | + |
| 32 | + |
| 33 | +def main(): |
| 34 | + heights = [6, 4, 8, 5, 2, 7] |
| 35 | + queries = [[0, 1], [0, 3], [2, 4], [3, 4], [2, 2]] |
| 36 | + assert Solution().leftmostBuildingQueries(heights, queries) == [2, 5, -1, 5, 2] |
| 37 | + |
| 38 | + heights = [5, 3, 8, 2, 6, 1, 4, 6] |
| 39 | + queries = [[0, 7], [3, 5], [5, 2], [3, 0], [1, 6]] |
| 40 | + assert Solution().leftmostBuildingQueries(heights, queries) == [7, 6, -1, 4, 6] |
| 41 | + |
| 42 | + |
| 43 | +if __name__ == '__main__': |
| 44 | + main() |
0 commit comments