Booking Concert Tickets in Groups LeetCode solution in Java

1
class BookMyShow {
2
  /** Segment tree class to store sum of a range and maximum available seats in a row */
3
  static class SegTree {
4
    long sum[]; // store sum of seats in a range
5
    long segTree[]; // store maximum seats in a range
6
    int m, n;
7

8
    public SegTree(int n, int m) {
9
      this.m = m;
10
      this.n = n;
11
      segTree = new long[4 * n];
12
      sum = new long[4 * n];
13
      build(0, 0, n - 1, m);
14
    }
15

16
    private void build(int index, int lo, int hi, long val) {
17
      if (lo == hi) {
18
        segTree[index] = val; // initialize segement tree with initial seat capacity
19
        sum[index] = val; // initialize "sum" with initial seat capacity of a row
20
        return;
21
      }
22
      int mid = (lo + hi) / 2;
23
      build(2 * index + 1, lo, mid, val); // build left sub tree
24
      build(2 * index + 2, mid + 1, hi, val); // build right sub tree
25
      segTree[index] =
26
          Math.max(
27
              segTree[2 * index + 1],
28
              segTree[2 * index + 2]); // maximum seats in a row for subtrees
29
      sum[index] = sum[2 * index + 1] + sum[2 * index + 2]; // sum of seats in a range
30
    }
31

32
    private void update(int index, int lo, int hi, int pos, int val) {
33
      /** Method to update segment tree based on the available seats in a row */
34
      if (lo == hi) {
35
        segTree[index] = val;
36
        sum[index] = val;
37
        return;
38
      }
39
      int mid = (lo + hi) / 2;
40
      if (pos <= mid) { // position to update is in left
41
        update(2 * index + 1, lo, mid, pos, val);
42
      } else { // position to update is in right
43
        update(2 * index + 2, mid + 1, hi, pos, val);
44
      }
45
      // update segment tree and "sum" based on the update in "pos" index
46
      segTree[index] = Math.max(segTree[2 * index + 1], segTree[2 * index + 2]);
47
      sum[index] = sum[2 * index + 1] + sum[2 * index + 2];
48
    }
49

50
    public void update(int pos, int val) {
51
      update(0, 0, n - 1, pos, val);
52
    }
53

54
    public int gatherQuery(int k, int maxRow) {
55
      return gatherQuery(0, 0, n - 1, k, maxRow);
56
    }
57

58
    private int gatherQuery(int index, int lo, int hi, int k, int maxRow) {
59
      /** Method to check if seats are available in a single row */
60
      if (segTree[index] < k || lo > maxRow) return -1;
61
      if (lo == hi) return lo;
62
      int mid = (lo + hi) / 2;
63
      int c = gatherQuery(2 * index + 1, lo, mid, k, maxRow);
64
      if (c == -1) {
65
        c = gatherQuery(2 * index + 2, mid + 1, hi, k, maxRow);
66
      }
67
      return c;
68
    }
69

70
    public long sumQuery(int k, int maxRow) {
71
      return sumQuery(0, 0, n - 1, k, maxRow);
72
    }
73

74
    private long sumQuery(int index, int lo, int hi, int l, int r) {
75
      if (lo > r || hi < l) return 0; // not in range
76
      if (lo >= l && hi <= r) return sum[index]; // in range
77
      int mid = (lo + hi) / 2;
78
      return sumQuery(2 * index + 1, lo, mid, l, r) + sumQuery(2 * index + 2, mid + 1, hi, l, r);
79
    }
80
  }
81

82
  SegTree segTree;
83
  int[] rowSeats; // stores avaiable seats in a row, helps to find the vacant seat in a row
84

85
  public BookMyShow(int n, int m) {
86
    segTree = new SegTree(n, m);
87
    rowSeats = new int[n];
88
    Arrays.fill(rowSeats, m); // initialize vacant seats count to "m" for all the rows
89
  }
90

91
  public int[] gather(int k, int maxRow) {
92
    int row = segTree.gatherQuery(k, maxRow); // find row which has k seats
93
    if (row == -1) return new int[] {}; // can't find a row with k seats
94
    int col = segTree.m - rowSeats[row]; // find column in the row which has k seats
95
    rowSeats[row] -= k; // reduce the seats
96
    segTree.update(row, rowSeats[row]); // update the segment tree
97
    return new int[] {row, col};
98
  }
99

100
  public boolean scatter(int k, int maxRow) {
101
    long sum = segTree.sumQuery(0, maxRow); // find the sum for the given range [0, maxRow]
102
    if (sum < k) return false; // can't find k seats in [0, maxRow]
103

104
    for (int i = 0; i <= maxRow && k != 0; i++) {
105
      if (rowSeats[i] > 0) { // if current row has seats then allocate those seats
106
        long t = Math.min(rowSeats[i], k);
107
        rowSeats[i] -= t;
108
        k -= t;
109
        segTree.update(i, rowSeats[i]); // update the segment tree
110
      }
111
    }
112
    return true;
113
  }
114
}