Count Subtrees With Max Distance Between Cities LeetCode solution in C++

1
class Solution {
2
 public:
3
  vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>> &edges) {
4
    // INT_MAX/2-1 used to avoid overflow
5
    vector<vector<int>> graph(n, vector<int>(n, INT_MAX / 2 - 1));
6
    for (auto e : edges) {
7
      int u = e[0] - 1, v = e[1] - 1;
8
      graph[u][v] = graph[v][u] = 1;
9
    }
10

11
    // floyed-warshall algorithm
12
    for (int k = 0; k < n; k++) {
13
      for (int i = 0; i < n; i++) {
14
        for (int j = 0; j < n; j++) {
15
          graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j]);
16
        }
17
      }
18
    }
19

20
    vector<int> res(n - 1, 0);
21

22
    /* Check for each combination of nodes and see if it's a subtree or not.
23
    If it's a subtree then see maximum distance that can be reached.*/
24

25
    // Each subtree
26
    // sub = subtree
27
    for (int sub = 0; sub < (1 << n); sub++) {
28
      // Number of nodes = nodes
29
      int nodes = __builtin_popcount(sub);
30
      // edges
31
      int edges = 0;
32
      // dist = distance
33
      int dist = 0;
34
      for (int i = 0; i < n; i++)
35
        if (sub & (1 << i)) {
36
          for (int j = i + 1; j < n; j++)
37
            if (sub & (1 << j)) {
38
              // if i to j edge increment e
39
              edges += graph[i][j] == 1;
40
              // maximum distance
41
              dist = max(dist, graph[i][j]);
42
            }
43
        }
44
      // chech its a subtree and distance > 0 if it's then its a subtree having
45
      // max d distance
46
      if (edges == nodes - 1 && dist > 0) {
47
        res[dist - 1]++;
48
      }
49
    }
50

51
    return res;
52
    // Please upvote
53
  }
54
};