ในข้อนี้ สามารถสรุปโจทย์ได้ดังนี้ กำหนดให้กราฟต้นไม้ที่มี $N$ โหนด และ $N - 1$ เส้นเชื่อม อยากทราบว่ามีคู่โหนด $u, v$ กี่คู่ที่จำนวนเส้นเชื่อมจาก $u$ ไป $v$ (หรือระยะห่างระหว่าง $u$ กับ $v$) มีค่าเท่ากับ $a_i$ สำหรับ $1 \leq i \leq k$

เนื่องจากกราฟเป็นต้นไม้ ดังนั้นเราสามารถคำนวณจำนวนคู่ $u, v$ สำหรับแต่ละ $a_i$ ด้วยการใช้ Dynamic Programming โดยเราจะเลือกโหนดใด ๆ มาหนึ่งโหนดเพื่อให้เป็นรากของต้นไม้ และกำหนดให้ $dp(u, i)$ เป็นจำนวนโหนด $v$ ใน Subtree ของ $u$ ที่ระยะห่างระหว่าง $v$ กับ $u$ เป็น $i$ สำหรับ Recurrence Formula จะเป็นดังนี้

$dp(u, i) = \sum \limits_{v} dp(v, i - 1)$ สำหรับ $1 \leq i \leq N - 1$ เมื่อ $v$ เป็นลูกของ $u$ ซึ่งมีที่มาจากการที่ระยะทางจาก $v$ มา $u$ เป็น $1$ ทำให้โหนดที่มีห่างจาก $v$ เท่ากับ $i - 1$ จะมีห่างจาก $u$ เท่ากับ $i$

นอกจากนี้ เราจะมีกรณีพิเศษคือ $dp(u, 0) = 1$ เพราะเราจะนับว่า $u$ ห่างจากตัวเองเป็นระยะ $0$ ด้วย

สำหรับการนับจำนวนคู่โหนดสำหรับค่า $a_i$ ใด ๆ เราสามารถทำได้โดย ขณะที่เรากำลังพิจารณาโหนด $u$ เราจะนับจำนวนคู่อันดับ $(v, w)$ ที่เส้นทางจาก $v$ มา $w$ ต้องผ่านโหนด $u$ ทำให้เราสามารถแบ่งเส้นทางจาก $v$ มา $w$ ออกเป็นเส้นทางจาก $v$ มา $u$ และ $u$ มา $w$ ตามลำดับ และเพื่อให้เส้นทางจาก $v$ มา $w$ ไม่ผ่านโหนดซ้ำ $v$ และ $w$ จะต้องเป็นโหนดในคนละ Subtree ของลูกของ $u$

ดังนั้น สำหรับค่า $a_i$ ใด ๆ จำนวนคู่โหนดที่ตรงตามเงื่อนไข และเส้นทางระหว่างโหนดต้องผ่าน $u$ จะเท่ากับ $\sum \limits_{j = 0}^{a_i - 2} dp(v, j) \cdot dp(w, a_i - j - 2)$ เมื่อ $v, w$ เป็นลูกของ $u$ ที่แตกต่างกัน สำหรับกรณีนี่เส้นทางผ่านโหนด $u$ ซึ่งมีที่มาจากการที่ (ระยะทางจากโหนดแรกของคู่อันดับมายัง $v$ $+$ ระยะทางจาก $v$ ไป $w$ ที่เท่ากับ $2$ เสมอ $+$ ระยะทางจาก $w$ ไปโหนดที่สองของคู่อันดับ) $= a_i$ และ $2 \cdot dp(u, a_i)$ สำหรับกรณีที่คู่อันดับมีตัวใดตัวหนึ่งเป็น $u$ โดยคูณ $2$ เพราะสามารถสลับตำแหน่งของคู่อันดับได้

หากทำตามวิธีด้านบนนี้ จะใช้เวลา $\mathcal{O}(N^2 K)$ แต่เราสามารถลดเวลาการทำงานลง ด้วยการเลือก $u$ ในการนับจำนวนคู่โหนดในลำดับที่เหมาะสมได้ โดย สมมติเราเลือก $u$ มา แล้วนับจำนวนคู่โหนดที่เส้นทางต้องผ่าน $u$ เป็นที่เรียบร้อยแล้ว เราจะลบ $u$ ออกจากต้นไม้ ทำให้กราฟกลายเป็นต้นไม้หลาย ๆ ต้น หลังจากนั้น เราจะทำการ Recursive ไปในต้นไม้แต่ละต้น เพื่อเลือก $u$ ถัดไปที่จะคำนวณ หากต้นไม้มีอยู่ $N$ โหนด และเราสามารถหาโหนด $u$ ที่เมื่อลบออกแล้ว จะทำให้ต้นไม้แต่ละต้นที่เหลือ มีจำนวนโหนดไม่เกิน $\lfloor \frac{N}{2} \rfloor$ เราจะสามารถลดเวลารวมของการ Recursive ลงเหลือเพียง $\mathcal{O}(N K \log N)$ ซึ่งโหนด $u$ ที่มีคุณสมบัตินี้ จะเรียกว่า Centroid ซึ่งสามารถพิสูจน์ได้ว่า สำหรับต้นไม้ใด ๆ จะมี Centroid อย่างน้อยหนึ่งโหนดเสมอ เทคนิคนี้มีชื่อเรียกว่า Centroid Decomposition

สำหรับการทำ Centroid Decomposition นั้น เราจะลดเวลาการคำนวณ $dp(u, i)$ ลง โดยแทนที่จะคิดจาก Recurrence Formula (ต้องคิดสำหรับทุก $u$ ใน Subtree) เราใช้การ DFS จาก $u$ เพื่อนับว่ามีกี่โหนดที่มีระยะห่างจาก $u$ ไป $1, 2, 3, ...$ เส้น (ไม่สนระยะห่างจากโหนดอื่นใน Subtree) โดยจัดการแยกแต่ละ Subtree ไม่ให้นับซ้ำ จึงกำหนดให้ $dp(i)$ เท่ากับจำนวนโหนดที่ห่างจาก $u$ เป็นระยะ $i$ และก่อนที่เราจะอัพเดทค่าของ $dp(i)$ สำหรับ Subtree ของ $v$ ซึ่งเป็นลูกของ $u$ เราจะทำการ DFS เพื่อนับจำนวนคู่อันดับที่มีตัวแรกเป็นโหนดใน Subtree ของ $v$ และอีกตัวมาจาก Subtree ก่อนหน้านี้ที่ได้อัพเดทค่าใน $dp(i)$ เป็นที่เรียบร้อยแล้ว หลังจากนั้นเราจะทำการ DFS อีกครั้งใน Subtree ของ $v$ เพื่ออัพเดทค่า $dp(i)$ แล้วคำนวณ Subtree ถัดไป หลังจากนั้น เมื่อพิจารณา Centroid $u$ เป็นที่เรียบร้อยแล้ว เราจะลบ $u$ ออกจากกราฟ แล้ว Recursive ไปยังลูกของ $u$ ที่ยังไม่โดนลบ เพื่อคำนวณคำตอบต่อไป วิธีนี้จะได้จำนวนคำตอบเพียงครึ่งเดียวจึงต้องคูณ $2$ ก่อนแสดงคำตอบ

โค้ดตัวอย่างดังนี้

#include <bits/stdc++.h>

using namespace std;

const int N = 1e5 + 5;

int n, k, A[N];
int sz[N], chk[N], book;
long long dp[N], t[N], ans[N];
vector<int> g[N];

int getsz(int u, int p) {
  sz[u] = 1;
  for (int v : g[u])
    if (v != p && !chk[v])
      sz[u] += getsz(v, u);
  return sz[u];
}

int find_cen(int u, int p, int all, pair<int, int> &ret) {
  int mx = all - sz[u];
  for (int v : g[u])
    if (v != p && !chk[v])
      mx = max(mx, find_cen(v, u, all, ret));
  ret = min(ret, pair<int, int>(mx, u));
  return sz[u];
}

void dfs(int u, int p, int len, bool fill) {
  if (!fill) {
    for (int i = 1; i <= k; i++)
      if (len <= A[i]) {
        if (len == A[i])
          ++ans[i];
        if (t[A[i] - len] == book)
          ans[i] += dp[A[i] - len];
      }
  } else {
    if (t[len] < book)
      dp[len] = 0;
    ++dp[len], t[len] = book;
  }
  for (int v : g[u])
    if (v != p && !chk[v])
      dfs(v, u, len + 1, fill);
}

void proc(int u) {
  if (getsz(u, u) <= 1)
    return;
  pair<int, int> ret(sz[u] + 1, -1);
  find_cen(u, u, sz[u], ret);
  u = ret.second;
  ++book;
  for (int v : g[u])
    if (!chk[v]) {
      dfs(v, u, 1, 0);
      dfs(v, u, 1, 1);
    }
  chk[u] = true;
  for (int v : g[u])
    if (!chk[v])
      proc(v);
}

int main() {
  scanf("%d %d", &n, &k);
  for (int i = 1; i <= k; i++)
    scanf("%d", A + i);
  for (int i = 1, a, b; i < n; i++) {
    scanf("%d %d", &a, &b);
    g[a].emplace_back(b), g[b].emplace_back(a);
  }
  proc(1);
  for (int i = 1; i <= k; i++)
    printf("%lld\n", ans[i] * 2ll);

  return 0;
}

สำหรับโค้ดด้านบนนี้ จะมีการใช้อาร์เรย์ t เนื่องจากเราจะต้องล้างค่าของอาร์เรย์ dp ทุก ๆ ครั้งที่คำนวณคำตอบของ Centroid อันใหม่

Time Complexity: $\mathcal{O}(N K \log N)$