CF 1611D. Weights Assignment For Tree Edges

Weights Assignment For Tree Edges

题意

给定 \(n\) 个结点的无根树和序列 \(b, p\) ，\(b_i\) 表示 \(i\) 结点的父结点，其中 \(b_{root} = root\) ，现在要给树上的每个边赋正权值，使得每个结点到根的距离 \(dist_i\) 满足 \(p\) 的排序。即如果在 \(p\) 数组中 \(i\) 点位置在 \(j\) 前面 ，则 \(dist_i < dist_j\) 。

分析

因为给边赋的权值一定为正数，那么对于一个结点，它的排名不可能比父结点靠前，也就是说 \(p\) 数组满足拓扑序。

对于 \(p = [3, 1, 2, 5, 4]\) ，构造 $dist_3 = 0(第一个一定是根节点), dist_1 = 2, dist_2 = 3 \ldots $ ，满足枚举的顺序，这个是结点到根的距离，只要记录一下父结点到跟的距离，相减就是这个边的权值。

Code

/* 终点是一切概率的结束，也是一切期望的开始 */
#include <bits/stdc++.h>
using namespace std;
const int N = 200010;

int b[N], p[N];
int ans[N], dis[N]; // ans为结点到父结点的边权，dis为结点到根节点的边权

void solve ()
{
    int n; cin >> n;
    for (int i = 1; i <= n; i ++ ) cin >> b[i];
    for (int i = 1; i <= n; i ++ ) cin >> p[i];
    for (int i = 1; i <= n; i ++ ) ans[i] = dis[i] = -1;
    ans[p[1]] = dis[p[1]] = 0; // 根节点为0
    for (int i = 2; i <= n; i ++ )
    {
        int j = p[i]; // 计算第i名
        if (dis[b[j]] == -1) return cout << "-1\n", void(); // 排名比父结点前
        ans[j] = i - dis[b[j]]; // 边权
        dis[j] = i; // 到root的距离
    }
    for (int i = 1; i <= n; i ++ ) cout << ans[i] << " \n"[i == n];
}

signed main ()
{
    cout.tie(0)->sync_with_stdio(0);
    int _; for (cin >> _; _--; ) solve();
    return 0;
}