Use this to learn the idea, then write your own version.
1struct SegmentTreeNode {2 int lo;3 int hi;4 std::unique_ptr<SegmentTreeNode> left;5 std::unique_ptr<SegmentTreeNode> right;6 int mx;7 long sum;8 SegmentTreeNode(int lo, int hi, std::unique_ptr<SegmentTreeNode>&& left,9 std::unique_ptr<SegmentTreeNode>&& right, int mx, long sum)10 : lo(lo),11 hi(hi),12 left(std::move(left)),13 right(std::move(right)),14 mx(mx),15 sum(sum) {}16};17 18class SegmentTree {19 public:20 explicit SegmentTree(int n, int m) : m(m), root(std::move(build(0, n - 1))) {}21 22 vector<int> maxRange(int k, int maxRow) {23 return maxRange(root, k, maxRow);24 }25 26 long sumRange(int maxRow) {27 return sumRange(root, 0, maxRow);28 }29 30 31 void substract(int row, int k) {32 substract(root, row, k);33 }34 35 private:36 const int m;37 std::unique_ptr<SegmentTreeNode> root;38 39 std::unique_ptr<SegmentTreeNode> build(int l, int r) {40 if (l == r)41 return make_unique<SegmentTreeNode>(l, r, nullptr, nullptr, m, m);42 const int mid = (l + r) / 2;43 std::unique_ptr<SegmentTreeNode> left = build(l, mid);44 std::unique_ptr<SegmentTreeNode> right = build(mid + 1, r);45 return make_unique<SegmentTreeNode>(l, r, std::move(left), std::move(right),46 max(left->mx, right->mx),47 left->sum + right->sum);48 }49 50 vector<int> maxRange(std::unique_ptr<SegmentTreeNode>& root, int k,51 int maxRow) {52 if (root->lo == root->hi) {53 if (root->sum < k || root->lo > maxRow)54 return {};55 return {root->lo, m - static_cast<int>(root->sum)}; 56 }57 58 if (root->left->mx >= k)59 return maxRange(root->left, k, maxRow);60 return maxRange(root->right, k, maxRow);61 }62 63 long sumRange(std::unique_ptr<SegmentTreeNode>& root, int i, int j) {64 if (root->lo == i && root->hi == j)65 return root->sum;66 const int mid = (root->lo + root->hi) / 2;67 if (j <= mid)68 return sumRange(root->left, i, j);69 if (i > mid)70 return sumRange(root->right, i, j);71 return sumRange(root->left, i, mid) + sumRange(root->right, mid + 1, j);72 }73 74 void substract(std::unique_ptr<SegmentTreeNode>& root, int row, int k) {75 if (root == nullptr)76 return;77 if (root->lo == root->hi && root->hi == row) {78 root->mx -= k;79 root->sum -= k;80 return;81 }82 const int mid = (root->lo + root->hi) / 2;83 if (row <= mid)84 substract(root->left, row, k);85 else86 substract(root->right, row, k);87 root->mx = max(root->left->mx, root->right->mx);88 root->sum = root->left->sum + root->right->sum;89 }90};91 92class BookMyShow {93 public:94 BookMyShow(int n, int m) : tree(n, m), seats(n, m) {}95 96 vector<int> gather(int k, int maxRow) {97 const vector<int> res = tree.maxRange(k, maxRow);98 if (res.size() == 2) {99 const int row = res[0];100 tree.substract(row, k);101 seats[row] -= k;102 }103 return res;104 }105 106 bool scatter(int k, int maxRow) {107 if (tree.sumRange(maxRow) < k)108 return false;109 110 while (k > 0)111 if (seats[minVacantRow] >= k) {112 tree.substract(minVacantRow, k);113 seats[minVacantRow] -= k;114 k = 0;115 } else {116 tree.substract(minVacantRow, seats[minVacantRow]);117 k -= seats[minVacantRow];118 seats[minVacantRow] = 0;119 ++minVacantRow;120 }121 122 return true;123 }124 125 private:126 SegmentTree tree;127 vector<int> seats; 128 int minVacantRow = 0;129};130