树上一些点的选 题解

题意简述

给你一棵 \(n\) 个节点以 \(1\) 为根的有根树，和一个整数 \(m\)。

对于树上每一个点 \(u\)，有三个权值 \(X, Y, Z\)。你需要在 \(u\) 的祖先里（不含 \(u\)）中选出至少 \(X\) 个点，记 \(S_1\) 表示这些点到 \(u\) 的距离之和；在 \(u\) 的后代里（不含 \(u\)）中选出至少 \(Y\) 个点，记 \(S_2\) 表示这些点到 \(u\) 的距离之和。选出的点数应恰好等于 \(m\)。对于每一个点，求 \(\min \vert S_1 - S_2 + Z \vert\)。如果没有合法的选点方案，输出 \(-1\)。

题目分析

这个距离和不喜欢，拆掉。

为方便叙述，记祖先中选出点集为 \(U\)，\(x = \vert U \vert\)，同理记后代中选出点集为 \(D\)，\(y = \vert D \vert\)。那么：

\[\begin{aligned} S_1 &= \sum _ {v \in U} \operatorname{dis}(v, u) \\ &= \sum _ {v \in U} \operatorname{dpt}(u) - \operatorname{dpt}(v) \\ &= x \operatorname{dpt}(u) - \sum _ {v \in U} \operatorname{dpt}(v) \\ \end{aligned} \]

同理 \(S_2 = -y \operatorname{dpt}(u) + \sum \limits _ {v \in D} \operatorname{dpt}(v)\)，那么所求值即为：

\[\begin{aligned} & \quad \ \min \left \vert \left (x \operatorname{dpt}(u) - \sum _ {v \in U} \operatorname{dpt}(v) \right ) - \left (-y \operatorname{dpt}(u) + \sum \limits _ {v \in D} \operatorname{dpt}(v) \right ) + Z \right \vert \\ &= \min \left \vert x \operatorname{dpt}(u) - \left ( \sum _ {v \in U} \operatorname{dpt}(v) \right ) + y \operatorname{dpt}(u) - \left ( \sum \limits _ {v \in D} \operatorname{dpt}(v) \right ) + Z \right \vert \\ &= \min \left \vert (x + y) \operatorname{dpt}(u) + Z - \left ( \sum _ {v \in U} \operatorname{dpt}(v) + \sum \limits _ {v \in D} \operatorname{dpt}(v) \right ) \right \vert \\ \end{aligned} \]

发现 \(x + y\) 为一定值 \(m\)，故而所求值即为两个 \(\sum\) 与一定值的最少差值绝对值，也就是尽量靠近那个定值。这就来到了此题难点所在。

一个很神的想法是去证明，这两个 \(\sum\) 的和的取值范围为一段连续区间，再扣掉常数个点，这样一个东西。

证明：

显然其有最值。从其取最大值的选点方案，往最小值调整。初始时所有点聚集在下方，取到最大值。

每次我们可以取一个 \(u\)，如果存在一个比它深度恰小 \(1\) 的点没被选，我们就可以让 \(u\) 变成那个深度比它小 \(1\) 的点，就将和往下调整了 \(1\)。

我们还有另一种方法。可以让 \(u\) 的一个孩子越过 \(u\)，走到 \(\operatorname{fa}(u)\)。但是这样导致将和往下调整了 \(2\)。注意到，我们可以选取另一个点，向下走一步，走到深度 \(+1\) 的位置来抵消。如果找不到这一个用来抵消的点，就不可以使用这一种方法，将和向下调整 \(1\)。如果此时使用前一种方法也不能下调 \(1\)，我们的取值范围就被挖去了这一个点，并且可以使用这一种方法向下调整 \(2\)，继续向最小值调整。

当前一种方法无法调整时，所有点都尽可能往上，\(u\) 一定有一个选中的孩子，我们也一定能使用这一种方法调整 \(2\)。

当所有点都不能走到深度 \(+1\) 的位置，其实等价于所有点都不能向其一个孩子移动。情况数不多，大力分讨调整了 \(2\) 后的情形：

1. 上方满了，下方为空。
2. 上方满了，下方无法调整。
3. 上方只有移动的点，下方为空。
4. 上方只有移动的点，下方无法调整。

接下来着重分析“下方无法调整”这种情况。比较好像的情况就是 \(u\) 的子树，除了一个孩子（我们用来越过 \(u\) 的点），都是满的。事实上，子树不满也有可能。此时说明 \(u\) 的部分孩子没被选择，也就说明了 \(u\) 的所有孩子都是叶子（否则可以来到深度 \(-1\) 的没被选择的 \(u\) 的孩子）。这种情况需要特判。

以上这 \(4\) 个大局面（\(6\) 个子局面）的 \(\sum + \sum + 1\) 是不能取到的。

那么，我们只需要求出这两个 \(\sum\) 的上下界即可。

对于深度之和的下界，就是尽可能多的点在 \(u\) 的祖先。说明此时有 \(x = x_{\max}\) 个点，深度分别为 \(1, \ldots, x_{\max}\)（\(\operatorname{dpt}(1) = 1\)）；剩下下方 \(y = m - x_{\max}\) 个点，尽可能靠上贴在 \(u\) 下方。

对于深度之和的上界，同理，尽可能地靠下。上方有 \(x = x_{\min}\) 个点，深度为 \(\operatorname{dpt}(u) - x_{\min}, \ldots, \operatorname{dpt}(u) - 1\)；下方 \(y = m - x_{\min}\) 个点，尽可能靠下。

求 \(x_{\min}, x_{\max}\) 是简单的。

\[\left\{\begin{matrix} X \leq x \leq \operatorname{dpt}(u) - 1 \\ Y \leq m - x \leq \operatorname{siz}(u) - 1 \end{matrix}\right. \Rightarrow \max \Big \{ X, m - \operatorname{siz}(u) + 1 \Big \} \leq x \leq \min \Big \{ \operatorname{dpt}(u) - 1, m - Y \Big \} \]

对于 \(1 \sim u\) 这条链上的答案很好统计，关键在于子树内，选出 \(k\) 个点，尽可能靠上 / 靠下的深度和。直接上线段树合并维护子树每个深度出现次数，加上前 / 后 \(k\) 大值的和查询。或者树上启发式合并，树状数组上二分。

前者是 \(\Theta(n \log n)\) 的。

代码

#include <cstdio>
#include <iostream>
#include <vector>
using namespace std;

const int N = 500010;
using lint = long long;

int n, m;
vector<int> edge[N];

struct Super_Sement_Tree {
	struct node {
		int lson, rson, cnt;
        lint sum;
	} tree[5000010];
	int tot, root[N];
    
    int & operator [] (int x) { return root[x]; }
    
    void modify(int &idx, int trl, int trr, int p, int v) {
        if (trl > p || trr < p) return;
        if (!idx) idx = ++tot;
        tree[idx].sum += v, ++tree[idx].cnt;
        if (trl == trr) return;
        int mid = (trl + trr) >> 1;
        modify(tree[idx].lson, trl, mid, p, v);
        modify(tree[idx].rson, mid + 1, trr, p, v);
    }
    
    int combine(int u, int v) {
        if (!u || !v) return u | v;
        tree[u].sum += tree[v].sum;
        tree[u].cnt += tree[v].cnt;
        tree[u].lson = combine(tree[u].lson, tree[v].lson);
        tree[u].rson = combine(tree[u].rson, tree[v].rson);
        return u;
    }
    
    lint query_pre(int idx, int trl, int trr, int k) {
        if (trl == trr) return 1ll * trl * k;
        int mid = (trl + trr) >> 1;
        if (k <= tree[tree[idx].lson].cnt) return query_pre(tree[idx].lson, trl, mid, k);
        return query_pre(tree[idx].rson, mid + 1, trr, k - tree[tree[idx].lson].cnt) + tree[tree[idx].lson].sum;
    }
    
    lint query_suf(int idx, int trl, int trr, int k) {
        if (trl == trr) return 1ll * trl * k;
        int mid = (trl + trr) >> 1;
        if (k <= tree[tree[idx].rson].cnt) return query_suf(tree[idx].rson, mid + 1, trr, k);
        return query_suf(tree[idx].lson, trl, mid, k - tree[tree[idx].rson].cnt) + tree[tree[idx].rson].sum;
    }
} yzh;

long long ans[N], sdpt[N];
int siz[N], dpt[N];
int X[N], Y[N], Z[N];

lint calc(lint l, lint r) {
    return 1ll * (l + r) * (r - l + 1) / 2;
}

void dfs(int now, int fa) {
    siz[now] = 1, dpt[now] = dpt[fa] + 1;
    bool all_leaf = true;
    for (int to: edge[now]) if (to != fa) {
        dfs(to, now), siz[now] += siz[to];
        yzh[now] = yzh.combine(yzh[now], yzh[to]);
        sdpt[now] += sdpt[to];
        all_leaf &= siz[to] == 1u;
    }
    int L = max(X[now], m - siz[now] + 1), R = min(m - Y[now], dpt[now] - 1);
    if (L > R) ans[now] = -1;
    else {
        long long val = Z[now] + 1ll * m * dpt[now];
        long long l = calc(1, R) + yzh.query_pre(yzh[now], 1, n, m - R);
        long long r = calc(dpt[now] - L, dpt[now] - 1) + yzh.query_suf(yzh[now], 1, n, m - L);
        if (val < l) ans[now] = l - val;
        else if (val > r) ans[now] = val - r;
        else {
            int son = (int)edge[now].size() - (fa != 0);
            ans[now] = 0;
            if (m == dpt[now] - 1 && val == calc(1, m) + 1) ans[now] = 1;
            if (m == dpt[now] - 1 + siz[now] - 2 && val == calc(1, dpt[now] - 1) + sdpt[now] - (dpt[now] + 1) + 1) ans[now] = 1;
            if (all_leaf && dpt[now] - 1 <= m && m <= dpt[now] - 1 + son - 1 && val == calc(1, dpt[now] - 1) + 1ll * (dpt[now] + 1) * (m - (dpt[now] - 1)) + 1) ans[now] = 1;
            if (m == 1 && val == (dpt[now] - 1) + 1) ans[now] = 1;
            if (m == 1 + siz[now] - 2 && val == (dpt[now] - 1) + sdpt[now] - (dpt[now] + 1) + 1) ans[now] = 1;
            if (all_leaf && 1 <= m && m <= 1 + son - 1 && val == (dpt[now] - 1) + 1ll * (dpt[now] + 1) * (m - 1) + 1) ans[now] = 1;
        }
    }
    sdpt[now] += dpt[now];
    yzh.modify(yzh[now], 1, n, dpt[now], dpt[now]);
}

signed main() {
    #ifndef XuYueming
    freopen("c.in", "r", stdin);
    freopen("c.out", "w", stdout);
    #endif
    scanf("%d%d", &n, &m);
    for (int i = 1, u, v; i < n; ++i) {
        scanf("%d%d", &u, &v);
        edge[u].emplace_back(v);
        edge[v].emplace_back(u);
    }
    for (int i = 1; i <= n; ++i) scanf("%d%d%d", &X[i], &Y[i], &Z[i]);
    dfs(1, 0);
    for (int i = 1; i <= n; ++i)
        printf("%lld ", ans[i]);
    return 0;
}