Count Pairs Of Nodes LeetCode solution in C++

1
class Solution {
2
  vector<int> st;
3

4
 public:
5
  void update(int tind, int tl, int tr, int ind, int val) {
6
    if (tl > tr) return;
7
    if (tl == tr) {
8
      st[tind] += val;
9
      return;
10
    }
11
    int tm = tl + ((tr - tl) >> 1), left = tind << 1;
12
    if (ind <= tm)
13
      update(left, tl, tm, ind, val);
14
    else
15
      update(left | 1, tm + 1, tr, ind, val);
16
    st[tind] = st[left] + st[left | 1];
17
  }
18

19
  int query(int tind, int tl, int tr, int ql, int qr) {
20
    if (tl > tr or qr < tl or ql > tr) return 0;
21
    if (ql <= tl and tr <= qr) return st[tind];
22
    int tm = tl + ((tr - tl) >> 1), left = tind << 1;
23
    return query(left, tl, tm, ql, qr) + query(left | 1, tm + 1, tr, ql, qr);
24
  }
25

26
  vector<int> countPairs(int n, vector<vector<int>> &a, vector<int> &q) {
27
    vector<int> f(n + 2, 0);
28
    vector<unordered_map<int, int>> mp(n + 2); // using a normal map gives TLE
29

30
    for (auto v : a) {
31
      int mx = max(v[0], v[1]), mn = min(v[1], v[0]);
32
      f[mx]++;
33
      f[mn]++;
34
      mp[mx][mn]++;
35
    }
36

37
    vector<int> ans(q.size(), 0);
38
    int m = a.size();
39
    st.resize(4 * m + 40, 0);
40

41
    for (int i = n; i; i--) {
42
      // reducing the common edges between current node and its adjacent
43
      for (auto e : mp[i]) {
44
        update(1, 0, m, f[e.first], -1);
45
        update(1, 0, m, f[e.first] - e.second, 1);
46
      }
47

48
      int curr = f[i];
49
      for (int j = 0; j < q.size(); j++) ans[j] += query(1, 0, m, max(0, q[j] - curr + 1), m);
50

51
      // reversing the changes made initially
52
      for (auto e : mp[i]) {
53
        update(1, 0, m, f[e.first], 1);
54
        update(1, 0, m, f[e.first] - e.second, -1);
55
      }
56

57
      update(1, 0, m, f[i], 1);
58
    }
59
    return ans;
60
  }
61
};