1617

1617. Count Subtrees With Max Distance Between Cities (Hard)

There are n cities numbered from 1 to n. You are given an array edges of size n-1, where edges[i] = [ui, vi] represents a bidirectional edge between cities ui and vi. There exists a unique path between each pair of cities. In other words, the cities form a tree.

A subtree is a subset of cities where every city is reachable from every other city in the subset, where the path between each pair passes through only the cities from the subset. Two subtrees are different if there is a city in one subtree that is not present in the other.

For each d from 1 to n-1, find the number of subtrees in which the maximum distance between any two cities in the subtree is equal to d.

Return an array of size n-1 where the dth element (1-indexed) is the number of subtrees in which the maximum distance between any two cities is equal to d.

Notice that the distance between the two cities is the number of edges in the path between them.

 

Example 1:

Input: n = 4, edges = [[1,2],[2,3],[2,4]]
Output: [3,4,0]
Explanation:
The subtrees with subsets {1,2}, {2,3} and {2,4} have a max distance of 1.
The subtrees with subsets {1,2,3}, {1,2,4}, {2,3,4} and {1,2,3,4} have a max distance of 2.
No subtree has two nodes where the max distance between them is 3.

Example 2:

Input: n = 2, edges = [[1,2]]
Output: [1]

Example 3:

Input: n = 3, edges = [[1,2],[2,3]]
Output: [2,1]

 

Constraints:

• 2 <= n <= 15
• edges.length == n-1
• edges[i].length == 2
• 1 <= ui, vi <= n
• All pairs (ui, vi) are distinct.

Related Topics:
Dynamic Programming

Similar Questions:

• Tree Diameter (Medium)

Solution 1.

We need to:

• traverse all the subset of nodes. We can use bitmask for this.
• check if the subset is connected. We can use DFS, BFS, UnionFind.
• get the maximum distance of the subtree. We can use Floyd algorithm to precompute the (minimal) distance between every node pairs. And in the subtree, just find the maximum distance by traversing every node pairs.

Time complexity:

• O(N^3) for Floyd
• O(2^N) for bitmask subset traversal
  • O(N) for memset, filling incl, and dfs.
  • O(N^2) for getting the maximum distance

So overall the time complexity is O(2^N * N^2).

// OJ: https://leetcode.com/problems/count-subtrees-with-max-distance-between-cities/
// Author: github.com/lzl124631x
// Time: O(2^N)
// Space: O(N^2)
class Solution {
    int seen[16] = {}, incl[16] = {}; // incl[i] == 1 if node `i` is included in mask
    vector<int> g[16];
    void dfs(int u, int &cnt) {
        if (!incl[u] || seen[u]) return;
        ++cnt;
        seen[u] = 1;
        for (int v : g[u]) dfs(v, cnt);
    }
public:
    vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>>& E) {
        int dist[16][16] = {};
        memset(dist, 0x3f, sizeof(dist));
        for (auto &e : E) {
            int u = e[0] - 1, v = e[1] - 1;
            g[u].push_back(v);
            g[v].push_back(u);
            dist[u][v] = dist[v][u] = 1;
        }
        for (int k = 0; k < n; ++k) { // use floyd to get the minimal distance between every two nodes.
            for (int i = 0; i < n; ++i) {
                for (int j = 0; j < n; ++j) {
                    dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j]);
                }
            }
        }
        vector<int> ans(n - 1);
        for (int mask = 1; mask < (1 << n); ++mask) {
            memset(seen, 0, sizeof(seen));
            memset(incl, 0, sizeof(incl));
            int cnt = __builtin_popcount(mask), start;
            if (cnt < 2) continue;
            for (int i = 0; i < n; ++i) {
                if (mask & (1 << i)) {
                    incl[i] = 1;
                    start = i;
                }
            }
            int c = 0;
            dfs(start, c); // count how many nodes we can visit.
            if (cnt != c) continue; // not all nodes are connected
            int mx = 0;
            for (int i = 0; i < n; ++i) {
                if (!incl[i]) continue;
                for (int j = i + 1; j < n; ++j) {
                    if (!incl[j]) continue;
                    mx = max(mx, dist[i][j]);
                }
            }
            ++ans[mx - 1];
        }
        return ans;
    }
};

Solution 2.

We can use two DFS to get the tree's diameter and by the way we can know if the subtree is connected by counting the dist.

// OJ: https://leetcode.com/problems/count-subtrees-with-max-distance-between-cities/
// Author: github.com/lzl124631x
// Time: O(2^N * N)
// Space: O(N + E)
class Solution {
    int incl[16] = {}, dist[16] = {}; // incl[i] == 1 if node `i` is included in mask
    vector<int> g[16];
    void dfs(int u, int &far) {
        for (int v : g[u]) {
            if (!incl[v] || dist[v]) continue;
            dist[v] = dist[u] + 1;
            if (dist[v] > dist[far]) far = v;
            dfs(v, far);
        }
    }
public:
    vector<int> countSubgraphsForEachDiameter(int n, vector<vector<int>>& E) {
        memset(dist, 0x3f, sizeof(dist));
        for (auto &e : E) {
            int u = e[0] - 1, v = e[1] - 1;
            g[u].push_back(v);
            g[v].push_back(u);
        }
        vector<int> ans(n - 1);
        for (int mask = 1; mask < (1 << n); ++mask) {
            int cnt = __builtin_popcount(mask), node, c = 0;
            if (cnt < 2) continue;
            memset(incl, 0, sizeof(incl));
            for (int i = 0; i < n; ++i) {
                if (mask & (1 << i)) {
                    incl[i] = 1;
                    node = i;
                }
            }
            memset(dist, 0, sizeof(dist));
            dist[node] = 1;
            dfs(node, node);
            for (int i = 0; i < n; ++i) c += incl[i] && dist[i];
            if (c != cnt) continue;
            memset(dist, 0, sizeof(dist));
            dist[node] = 1;
            dfs(node, node);
            ++ans[dist[node] - 2];
        }
        return ans;
    }
};