|
| 1 | +class SegmentTree: |
| 2 | + def __init__(self, baskets: list[int]): |
| 3 | + self.baskets = baskets |
| 4 | + self.m = len(baskets) |
| 5 | + self.tree = [0] * (4 * self.m) |
| 6 | + self.build(0, 0, self.m - 1) |
| 7 | + |
| 8 | + def build(self, curr: int, left: int, right: int) -> None: |
| 9 | + if left == right: |
| 10 | + self.tree[curr] = self.baskets[left] |
| 11 | + return |
| 12 | + mid = (left + right) // 2 |
| 13 | + self.build(curr * 2 + 1, left, mid) |
| 14 | + self.build(curr * 2 + 2, mid + 1, right) |
| 15 | + self.tree[curr] = max(self.tree[curr * 2 + 1], self.tree[curr * 2 + 2]) |
| 16 | + |
| 17 | + def update(self, curr: int, left: int, right: int, idx: int) -> None: |
| 18 | + if left == right: |
| 19 | + self.tree[curr] = 0 |
| 20 | + return |
| 21 | + mid = (left + right) // 2 |
| 22 | + if idx <= mid: |
| 23 | + self.update(curr * 2 + 1, left, mid, idx) |
| 24 | + else: |
| 25 | + self.update(curr * 2 + 2, mid + 1, right, idx) |
| 26 | + self.tree[curr] = max(self.tree[curr * 2 + 1], self.tree[curr * 2 + 2]) |
| 27 | + |
| 28 | + def query(self, curr: int, left: int, right: int, x: int) -> int: |
| 29 | + if self.tree[curr] < x: |
| 30 | + return -1 |
| 31 | + if left == right: |
| 32 | + return left |
| 33 | + mid = (left + right) // 2 |
| 34 | + if self.tree[curr * 2 + 1] >= x: |
| 35 | + return self.query(curr * 2 + 1, left, mid, x) |
| 36 | + return self.query(curr * 2 + 2, mid + 1, right, x) |
| 37 | + |
| 38 | + |
| 39 | +class Solution: |
| 40 | + def numOfUnplacedFruits(self, fruits: list[int], baskets: list[int]) -> int: |
| 41 | + segment_tree = SegmentTree(baskets) |
| 42 | + unplaced = 0 |
| 43 | + for fruit in fruits: |
| 44 | + idx = segment_tree.query(0, 0, segment_tree.m - 1, fruit) |
| 45 | + if idx == -1: |
| 46 | + unplaced += 1 |
| 47 | + else: |
| 48 | + segment_tree.update(0, 0, segment_tree.m - 1, idx) |
| 49 | + return unplaced |
| 50 | + |
| 51 | + |
| 52 | +def main(): |
| 53 | + fruits = [4, 2, 5] |
| 54 | + baskets = [3, 5, 4] |
| 55 | + assert Solution().numOfUnplacedFruits(fruits, baskets) == 1 |
| 56 | + |
| 57 | + fruits = [3, 6, 1] |
| 58 | + baskets = [6, 4, 7] |
| 59 | + assert Solution().numOfUnplacedFruits(fruits, baskets) == 0 |
| 60 | + |
| 61 | + |
| 62 | +if __name__ == '__main__': |
| 63 | + main() |
0 commit comments