Count Pairs Of Nodes LeetCode solution in Java

1
class Solution {
2
  // Time  :  O(E + N + Q)
3
  // Space :  O(N + E)
4
  public int[] countPairs(int n, int[][] edges, int[] queries) {
5
    // Key: edge ID   Value: count of duplicate edge
6
    Map<Integer, Integer> edgeCount = new HashMap<>();
7
    // degree[i] : number of edge for node i (0-indexed)
8
    int[] degree = new int[n];
9

10
    for (int[] e : edges) { // recording all edges  ==>  O(E) Time / Space
11
      int u = e[0] - 1;
12
      int v = e[1] - 1;
13
      degree[u]++;
14
      degree[v]++;
15

16
      int eId = Math.min(u, v) * n + Math.max(u, v);
17
      edgeCount.put(eId, edgeCount.getOrDefault(eId, 0) + 1);
18
    }
19

20
    // Key: Degree   Value: Frequency (number of nodes with that degree)
21
    Map<Integer, Integer> degreeCount = new HashMap<>(); //   ==> O(U) Time / Space
22
    int maxDegree = 0;
23
    for (int d : degree) {
24
      degreeCount.put(d, degreeCount.getOrDefault(d, 0) + 1);
25
      maxDegree = Math.max(maxDegree, d);
26
    }
27

28
    int[] count = new int[2 * maxDegree + 1];
29

30
    for (int d1 : degreeCount.keySet()) { // O(E)-time  (seems to be O(U ^ 2))
31
      for (int d2 : degreeCount.keySet()) {
32
        count[d1 + d2] +=
33
            (d1 == d2)
34
                ? degreeCount.get(d1) * (degreeCount.get(d1) - 1)
35
                : degreeCount.get(d1) * degreeCount.get(d2);
36
      }
37
    }
38
    for (int i = 0; i < count.length; i++) count[i] /= 2; // each pair is counted twice
39

40
    for (int e : edgeCount.keySet()) { // O(E)-time
41
      int u = e / n;
42
      int v = e % n;
43

44
      count[degree[u] + degree[v]]--;
45
      count[degree[u] + degree[v] - edgeCount.get(e)]++;
46
    }
47

48
    for (int i = count.length - 2; i >= 0; i--) { // O(U)-time
49
      count[i] += count[i + 1];
50
    }
51

52
    int[] res = new int[queries.length]; // O(Q)-time
53
    for (int q = 0; q < queries.length; q++) {
54
      res[q] = ((queries[q] + 1) >= count.length) ? 0 : count[queries[q] + 1];
55
    }
56
    return res;
57
  }
58
}