翻转树边(树形dp)

4381. 翻转树边

翻转树边

思路

样例输入：

8
1 2
3 2
4 3
4 5
6 5
6 7
8 7

样例输出：

3
4 6 8

代码：

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;

const int N = 4e5 + 10, INF = 0x3f3f3f3f;
int e[N], ne[N], w[N], h[N], idx;
int n;
int down[N], up[N];
void add(int a, int b, int c)
{
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx ++;
}

void dfs1(int u, int fa)
{
    for(int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if(fa == j) continue;
        if(!w[i]) up[j] = up[u] + 1 + down[u] - down[j];
        else up[j] = up[u] + down[u] - 1 - down[j];
        dfs1(j, u);
    }
    return;
}

void dfs2(int u, int fa)
{
    int sum = 0;
    for(int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if(fa == j) continue;
        dfs2(j, u);
        if(w[i]) sum ++;
        sum += down[j];
    }
    down[u] = sum;
    return;
}

int main()
{
    scanf("%d", &n);
    
    memset(h, -1, sizeof h);
    
    for(int i = 1; i < n; i ++ )
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b, 0), add(b, a, 1);
    }
    
    //每个节点往下走需要翻的边
    dfs2(1, -1);
    
    //每个节点往上需要翻的边
    dfs1(1, -1);
    
    //求最小值
    int res = INF;
    for(int i = 1; i <= n; i ++ ) res = min(res, down[i] + up[i]);
    cout << res << endl;
    
    //输出所有满足最小条件的边
    for(int i = 1; i <= n; i ++ )
        if(res == down[i] + up[i])  cout << i << " ";
    
    return 0;
}