https://leetcode.com/problems/booking-concert-tickets-in-groups/
如何把一个问题转化成用线段树解决的问题，这个题是一个很好的示范。

个人解答

  1. class Node:
  2.     def __init__(self, v = 0):
  3.         self.left = None
  4.         self.right = None
  5.         self.max = v
  6.         self.sum = v
  8. class SegTree:
  9.     def __init__(self, n, initVal):
  10.         self.initVal = initVal
  11.         self.n = n
  12.         self.tree = self._build(0, n - 1, initVal)
  14.     def _build(self, l, r, initVal):
  15.         if l == r:
  16.             return Node(initVal)
  17.         node = Node()
  18.         mid = (l + r) // 2
  19.         node.left = self._build(l, mid, initVal)
  20.         node.right = self._build(mid + 1, r, initVal)
  21.         node.max = max(node.left.max, node.right.max)
  22.         node.sum = node.left.sum + node.right.sum
  23.         return node
  25.     def _decrease(self, node, l, r, i, v):
  26.         if l == r:
  27.             node.max -= v
  28.             node.sum -= v
  29.             return
  30.         mid = (l + r) // 2
  31.         if i <= mid:
  32.             self._decrease(node.left, l, mid, i, v)
  33.         else:
  34.             self._decrease(node.right, mid + 1, r, i, v)
  35.         node.max = max(node.left.max, node.right.max)
  36.         node.sum = node.left.sum + node.right.sum
  38.     def decrease(self, i, v):
  39.         self._decrease(self.tree, 0, self.n - 1, i, v)
  41.     def _queryAvaliable(self, node, l, r, target, maxRow):
  42.         if l > maxRow or node.max < target:
  43.             return []
  44.         if l == r:
  45.             # 要返回所在行以及当前行第一个空的index
  46.             return [l, self.initVal - node.max]
  47.         mid = (l + r) // 2
  48.         if node.left.max >= target:
  49.             return self._queryAvaliable(node.left, l, mid, target, maxRow)
  50.         return self._queryAvaliable(node.right, mid + 1, r, target, maxRow)
  52.     def queryAvaliable(self, target, maxRow):
  53.         return self._queryAvaliable(self.tree, 0, self.n - 1, target, maxRow)
  55.     def _querySum(self, node, l, r, maxRow):
  56.         if l > maxRow:
  57.             return 0
  58.         # 区间小于要请求的区间，直接返回！
  59.         if r <= maxRow or l == r:
  60.             return node.sum
  61.         mid = (l + r) // 2
  62.         return self._querySum(node.left, l, mid, maxRow) + self._querySum(node.right, mid + 1, r, maxRow)
  64.     def querySum(self, maxRow):
  65.         return self._querySum(self.tree, 0, self.n - 1, maxRow)
  68. class BookMyShow:
  70.     def __init__(self, n: int, m: int):
  71.         self.tree = SegTree(n, m)
  72.         # 记录每行剩下多少，以及第一个remian > 0的行，用来scatter的时候，选取合适的位置更新
  73.         self.remain = [m] * n
  74.         self.n = n
  75.         self.remainIdx = 0
  77.     def gather(self, k: int, maxRow: int) -> List[int]:
  78.         res = self.tree.queryAvaliable(k, maxRow)
  79.         if res and res[0] >= 0:
  80.             self.tree.decrease(res[0], k)
  81.             self.remain[res[0]] -= k
  82.         return res
  85.     def scatter(self, k: int, maxRow: int) -> bool:
  86.         if self.tree.querySum(maxRow) < k:
  87.             return False
  88.         # 需要从最小的地方开始更新
  89.         i = self.remainIdx
  90.         while k:
  91.             if k >= self.remain[i]:
  92.                 # remain为0
  93.                 self.tree.decrease(i, self.remain[i])
  94.                 k -= self.remain[i]
  95.                 self.remain[i] = 0
  96.                 i += 1
  97.             else:
  98.                 self.tree.decrease(i, k)
  99.                 self.remain[i] -= k
  100.                 k = 0
  101.         self.remainIdx = i
  103.         return True
  107. # Your BookMyShow object will be instantiated and called as such:
  108. # obj = BookMyShow(n, m)
  109. # param_1 = obj.gather(k,maxRow)
  110. # param_2 = obj.scatter(k,maxRow)

题目分析

问题转化

本题看似是两个互不相干的操作，并不是常见的线段树查询/修改的典型模式。但仔细一看，不管是scatter还是gatter，都对应了两个行为：

  • 查询一个符合条件的区间
  • 修改一个/多个值

涉及到反复的修改和查询，线段树很适合了。
再想一想怎么利用区间状态：

  • gather需要一个连续的空间，其实我们要找到剩余空间大于给定值的行，那么区间信息可以存区间里的最大remain，这样可以快速提前拒绝
  • cluster可以分散，但题目的要求可以得知，分配是连续的，也就是连续区间的sum满足即可，那么区间也需要记录sum值

至此区间的值都记录好，接下来的修改和查询就比较自然了。

线段树

线段树的实现，其实不用想着太多的实现trick，就记住核心的思想，然后老老实实用指针实现就可以。

  • 叶节点存具体的值，根节点存区间的值，不管是max/min/sum
  • 每次查询/更新都是从root开始，二分处理
  • 注意提前返回，如querySum的时候，根据区间判断是否可以直接返回