检查是否可以将树拆分为 K 个相等的连通分量

简介

给定一棵树，需要检查是否可以将它拆分为 K 个相等的连通分量。连通分量即为在树中无法互相到达的子树。

方法

方法一：深度优先搜索

通过深度优先搜索，可以遍历整棵树。我们可以记录每个节点的子树大小，这样在搜索时就可以得到每棵子树的大小。接下来，我们可以将这些子树大小排序，然后从最大的开始选，看它是否能够作为一个连通分量的一部分。如果不能，就继续选择之后的子树进行尝试，直到所有的子树都被分配到某一个连通分量中，或者分配的连通分量的数量已经达到了 K。

时间复杂度为 O(n log n)。

方法二：二分答案

在这个方法中，我们可以二分答案。我们可以从 1 到所有节点权值的和的平均数（即设所有节点权值之和为 sum，节点数量为 n，则均值为 sum/n）进行二分。假设当前二分的值为 mid，则我们可以检查是否存在一种连通分量的划分方案，使得每个连通分量的权值之和小于等于 mid。

为了检查这一点，我们可以使用深度优先搜索遍历整棵树，然后对于每个连通分量，统计其权值之和。如果存在一种方案，使得每个连通分量的权值之和小于等于 mid，则必须继续缩小 mid。如果不存在这样的方案，则必须增加 mid。

时间复杂度为 O(n log^2 n)。

示例代码

方法一

bool check_equal_partitions_DFS(TreeNode* root, int k) {
    vector<int> node_counts;
    function<int(TreeNode*)> dfs = [&](TreeNode* node) {
        if (!node) return 0;
        int count = 1;
        for (auto child : node->children) {
            count += dfs(child);
        }
        node_counts.push_back(count);
        return count;
    };
    dfs(root);
    sort(node_counts.rbegin(), node_counts.rend());
    return search(node_counts, 0, k, 0);
}

bool search(const vector<int>& node_counts, int start, int k, int target) {
    if (k == 0) return true;
    if (target == 0) return search(node_counts, 0, k - 1, target);
    for (int i = start; i < node_counts.size(); i++) {
        if (node_counts[i] <= target) {
            if (search(node_counts, i + 1, k, target - node_counts[i])) {
                return true;
            }
        }
    }
    return false;
}

方法二

bool check_equal_partitions_binary_search(TreeNode* root, int k) {
    vector<int> node_values;
    function<int(TreeNode*)> dfs = [&](TreeNode* node) {
        if (!node) return 0;
        int sum = node->val;
        for (auto child : node->children) {
            sum += dfs(child);
        }
        node_values.push_back(sum);
        return sum;
    };
    dfs(root);
    int sum = accumulate(node_values.begin(), node_values.end(), 0);
    int low = 1, high = sum / k;
    while (low <= high) {
        int mid = low + (high - low) / 2;
        if (can_divide(node_values, k, mid)) {
            high = mid - 1;
        } else {
            low = mid + 1;
        }
    }
    return low > sum / k;
}

bool can_divide(const vector<int>& node_values, int k, int max_sum) {
    vector<int> sums(k, 0);
    return divide(node_values, sums, 0, max_sum);
}

bool divide(const vector<int>& node_values, vector<int>& sums, int i, int max_sum) {
    if (i == node_values.size()) {
        return true;
    }
    for (int j = 0; j < sums.size(); j++) {
        if (sums[j] + node_values[i] <= max_sum) {
            sums[j] += node_values[i];
            if (divide(node_values, sums, i + 1, max_sum)) {
                return true;
            }
            sums[j] -= node_values[i];
        }
        if (sums[j] == 0) {
            break;
        }
    }
    return false;
}