Wbtree - histcat

Problem#

Given a rooted tree where each node is either black or white. Node $1$ is the root.

For each white node, find a black node in its subtree to match it, where each black node can only match with one white node. You need to find the minimum sum of distances between all white nodes and their paired black nodes.

The distance between two nodes in the tree is defined as the number of edges on the simple path between them.

It is guaranteed that a solution exists.

$1 < n < 10^6$

Solution#

First, we can consider a greedy approach.

If we take white nodes as the main focus, we can enumerate the white nodes from deep to shallow (since deeper white nodes have fewer black nodes to choose from), and for each white node, maintain the closest black node in its subtree for pairing.

However, this approach seems to require a union-find data structure, which I am not familiar with, todo.

So we consider taking black nodes as the main focus, enumerating black nodes from shallow to deep (since shallower black nodes have fewer choices), and for each black node, find the deepest unpaired white node on the path to the root, which can be maintained using a union-find structure, see the code.

Code#

#include<bits/stdc++.h>

using namespace std;

//----------------------------------------
const int N = 1e6 + 10;
int n;
bool color[N];
int fa[N], depth[N];
int head[N], nxt[N << 1], to[N << 1], cnt;
int vis[N];
long long white_dis_sum;
int vis_of_p[N];
void add(int x, int y)
{
    to[++cnt] = y;
    nxt[cnt] = head[x];
    head[x] = cnt;
}

void dfs(int u, int anc, int first_white)
{
    depth[u] = depth[anc] + 1;
    fa[u] = first_white;
    if(color[u] == 0) white_dis_sum += depth[u];
    for(int i = head[u]; i; i = nxt[i])
    {
        int v = to[i];
        if(v == anc) continue;
        if(color[u]/*black*/)
        {
            dfs(v, u, first_white);
        }
        else/*white*/
        {
            dfs(v, u, u);
        }
    }
}

int findfa_with_vis_equals_to_0(int u)
{
//	cout << "qaq" << u << endl;
    if(u == -1 || u == 0x3f3f3f3f) return 0x3f3f3f3f;
    if(vis[u] == 0) return u;
    return fa[u] = findfa_with_vis_equals_to_0(fa[u]);
}

int main()
{
    // for(int i = 1;i <= n;i++)
    // {
    //     fa[i] = i;
    // }
//    clock_t c1 = clock();
// #ifdef LOCAL
    freopen("wbtree.in", "r", stdin);
    freopen("wbtree.out", "w", stdout);
// #endif
//----------------------------------------
    ios::sync_with_stdio(0), cin.tie(0);
    cin >> n;
    for(int i = 1;i <= n;i++)
        cin >> color[i];
    int u, v;
    for(int i = 1;i <= n - 1;i++)
    {
        cin >> u >> v;
        add(u, v), add(v, u);
    }
    // If not found, return -1
    dfs(1, 0, -1);

    queue <int> qwq;
    qwq.push(1);
    long long ans = 0;
//    for(int i = 1;i <= n;i++)
//    {
//    	cout << i << ": " << fa[i] << endl;
//	}
    while(qwq.size())
    {
        int t = qwq.front();
//    	cout << "qwq"<<t << endl; 
		vis_of_p[t] = 1;
        for(int i = head[t]; i; i = nxt[i])
        {
            int v = to[i];
            if(vis_of_p[v]) continue;
            qwq.push(v);
        }
        qwq.pop();
        if(color[t] == 0) continue;
        int deepest_white = findfa_with_vis_equals_to_0(fa[t]);
//        cout << t << "的deepest_white: " << deepest_white << endl; 
        if(deepest_white == 0x3f3f3f3f) continue;
        vis[deepest_white] = 1;
        ans += depth[t];
    }
    cout <<  ans - white_dis_sum;
//----------------------------------------
//end:
//    cerr << "Time Used:" << clock() - c1 << "ms" << endl;
    return 0;
}

/*
 * ┌───┐   ┌───┬───┬───┬───┐ ┌───┬───┬───┬───┐ ┌───┬───┬───┬───┐ ┌───┬───┬───┐
 * │Esc│   │ F1│ F2│ F3│ F4│ │ F5│ F6│ F7│ F8│ │ F9│F10│F11│F12│ │P/S│S L│P/B│  ┌┐    ┌┐    ┌┐
 * └───┘   └───┴───┴───┴───┘ └───┴───┴───┴───┘ └───┴───┴───┴───┘ └───┴───┴───┘  └┘    └┘    └┘
 * ┌───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───────┐ ┌───┬───┬───┐ ┌───┬───┬───┬───┐
 * │~ `│! 1│@ 2│# 3│$ 4│% 5│^ 6│& 7│* 8│( 9│) 0│_ -│+ =│ BacSp │ │Ins│Hom│PUp│ │Num│ / │ * │ - │
 * ├───┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─────┤ ├───┼───┼───┤ ├───┼───┼───┼───┤
 * │ Tab │ Q │ W │ E │ R │ T │ Y │ U │ I │ O │ P │{ [│} ]│ | \ │ │Del│End│PDn│ │ 7 │ 8 │ 9 │   │
 * ├─────┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴─────┤ └───┴───┴───┘ ├───┼───┼───┤ + │
 * │ Caps │ A │ S │ D │ F │ G │ H │ J │ K │ L │: ;│" '│ Enter  │               │ 4 │ 5 │ 6 │   │
 * ├──────┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴────────┤     ┌───┐     ├───┼───┼───┼───┤
 * │ Shift  │ Z │ X │ C │ V │ B │ N │ M │< ,│> .│? /│  Shift   │     │ ↑ │     │ 1 │ 2 │ 3 │   │
 * ├─────┬──┴─┬─┴──┬┴───┴───┴───┴───┴───┴──┬┴───┼───┴┬────┬────┤ ┌───┼───┼───┐ ├───┴───┼───┤ E││
 * │ Ctrl│ Win│ Alt│         Space         │ Alt│ Win│Menu│Ctrl│ │ ← │ ↓ │ → │ │   0   │ . │←─┘│
 * └─────┴────┴────┴───────────────────────┴────┴────┴────┴────┘ └───┴───┴───┘ └───────┴───┴───┘
 */