2023河南萌新联赛第（二）场：河南工业大学 F - 最短距离

题目描述

给定一棵包含 $n$ 个顶点的树 $T$ ，以及 $m$ 个查询请求。每个查询包含三个参数 $:x、y$ 和 $k$。其中 $x$ 和 $y$ 是树中的两个顶点，$k$ 是一个整数。对于每个查询，你需要计算树中所有顶点到从 $x$ 到 $y$ 的简单路径上的最近距离恰好为 $k$ 的顶点数量。

输入描述:

第一行，两个正整数n和m。 $（1\le n,m\le 1e5）$
接下来 $n-1$ 行，每行两个正整数 $x,y$ 。点 $x$ 和点 $y$ 之间有一条边。$（1\le x,y\le n)$
最后 $m$ 行，每行三个正整数$x，y，k$。$（1\le x,y\le n，1\le k\le 100）$

输出描述:

对于每次询问，输出一个整数作为答案
每个答案占一行

示例1

输入

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

输出

2
1

由于 $k$ 较小，我们可以先用一个简单的树形 DP 预处理出每个节点的子树中距离该节点距离为 $0-k$ 的节点数量，
记为 $dp[n][k]$。
然后再进行一次换根 DP，计算每个节点的非子树节点距离该节点的距离为 $0-k$ 的节点数量。同时分别给子树中
距离该节点的距离为 $0-k$ 的 $dp$ 数组做一个根节点到该位置路径上的前缀和。
对于每个查询 ：

1. 先计算 $x$ 和 $y$的最近公共祖先 LCA。
2. 在以 LCA 为根的子树中，到达 $x-y$ 简单路径距离为 $d$ 的节点数量
   pre[x][k] + pre[y][k] - 2 * pre[w][k] + dp[w][k]) - (pre[x][k - 1] + pre[y][k - 1] - 2 * pre[w][k - 1]) + f[w][k]
   （通过前缀和计算，在以一个节点为根的子树中，到该节点距离为 $d-1$ 的节点到其父节点的距离为 $d$ ，但是这些节点并不符合题意，故减去）。
3. 再加上距离 LCA 为 的非子树节点数量。
4. 总的结果就是上面两项相加。

总的时间复杂度：
预处理树形 DP 和换根 DP 均为 $O(nk)$ 。
查询时，计算最近公共祖先为 $O(logn)$ ，计算结果为 $O(1)$ 。
共有 $q$ 次查询。
故总的时间复杂度为 $O(n\times k+q\times logn)$ 。

import java.io.*;
import java.util.ArrayList;public class Main {static int N = 100010;static ArrayList[] e = new ArrayList[N];static int[][] dp = new int[N][110];static int[][] f = new int[N][110];static long[][] pre = new long[N][110];static int[] dep = new int[N];static int[][] fa = new int[N][21];public static void dfs_1(int u, int fr) {dp[u][0] = 1;dep[u] = dep[fr] + 1;fa[u][0] = fr;for (int i = 1; i <= 20; i++) {fa[u][i] = fa[fa[u][i - 1]][i - 1];}for (Object v : e[u]) {int j = (int) v;if (j == fr) continue;dfs_1(j, u);for (int i = 1; i <= 100; i++)dp[u][i] += dp[j][i - 1];}}public static void dfs_2(int u, int fr) {for (Object v : e[u]) {int j = (int) v;if (j == fr) continue;for (int i = 1; i <= 100; i++) {f[j][i] = dp[u][i - 1] + f[u][i - 1];if (i >= 2) f[j][i] -= dp[j][i - 2];}dfs_2(j, u);}}public static void dfs_3(int u, int fr) {for (int i = 0; i <= 100; i++) {pre[u][i] = dp[u][i] + pre[fr][i];}for (Object v : e[u]) {int j = (int) v;if (j == fr) continue;dfs_3(j, u);}}public static int LCA(int u, int v) {if (dep[u] < dep[v]) {int t = u;u = v;v = t;}int k = dep[u] - dep[v];for (int i = 0; i <= 20; i++) {if ((k & (1 << i)) != 0) {u = fa[u][i];}}if (u == v) return u;for (int i = 20; i >= 0; i--) {if (fa[u][i] != fa[v][i]) {u = fa[u][i];v = fa[v][i];}}return fa[u][0];}public static void main(String[] args) throws IOException {BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));String[] str = bf.readLine().split(" ");int n = Integer.parseInt(str[0]);int m = Integer.parseInt(str[1]);for (int i = 0; i <= n; i++) {e[i] = new ArrayList();}for (int i = 1; i <= n - 1; i++) {str = bf.readLine().split(" ");int x = Integer.parseInt(str[0]);int y = Integer.parseInt(str[1]);e[x].add(y);e[y].add(x);}dfs_1(1, 0);dfs_2(1, 0);dfs_3(1, 0);while (m-- > 0) {str = bf.readLine().split(" ");int x = Integer.parseInt(str[0]);int y = Integer.parseInt(str[1]);int k = Integer.parseInt(str[2]);int w = LCA(x, y);bw.write((pre[x][k] + pre[y][k] - 2 * pre[w][k] + dp[w][k]) - (pre[x][k - 1] + pre[y][k - 1] - 2 * pre[w][k - 1]) + f[w][k] + "\n");}bw.close();}
}