Học nhóm (Đề tự luyện HSG Tin quốc gia – Bài 3) – Phân rã trọng tâm cây (centroid) và mẹo “chỉ bạn có r nhỏ hơn mới quyết định”: sắp xếp theo r + cây Fenwick
Đề bài
Học nhóm. Lớp học thuật toán của thầy Sơn ở xã vườn Phong Điền (Cần Thơ) có $N$ học trò; nhà của bạn thứ $i$ gọi là nhà $i$. Các nhà được nối bởi $N - 1$ con đường đal: con đường thứ $j$ nối nhà $u_j$ và $v_j$, dài $w_j$ mét, và giữa hai nhà bất kỳ có đúng một đường đi (các nhà và con đường tạo thành một cây). Gọi $d(u, v)$ là độ dài đường đi giữa hai nhà.
Thầy Sơn muốn ghép các cặp học nhóm: hai bạn thay phiên nhau đạp xe sang nhà bạn kia. Bạn thứ $i$ chỉ được đạp đi xa nhà không quá $r_i$ mét. Vì vậy hai bạn $u \ne v$ ghép được thành cặp khi và chỉ khi cả hai đều tới được nhà bạn kia: $d(u, v) \le r_u$ và $d(u, v) \le r_v$. Hãy đếm số cặp $\{u, v\}$ ghép được.
Dữ liệu vào: dòng đầu là $N$ $(1 \le N \le 7 \cdot 10^4)$; dòng hai là $r_1, \ldots, r_N$ $(0 \le r_i \le 10^{15})$; $N - 1$ dòng sau, mỗi dòng ba số $u_j, v_j, w_j$ $(1 \le w_j \le 10^9)$.
Kết quả: số cặp học nhóm có thể ghép.
- Subtask 1 (20%): $N \le 2000$.
- Subtask 2 (20%): các nhà nằm trên một đường thẳng, con đường thứ $j$ nối nhà $j$ và nhà $j + 1$.
- Subtask 3 (25%): $r_1 = r_2 = \ldots = r_N$.
- Subtask 4 (35%): không có ràng buộc gì thêm.
Giới hạn thời gian: C++ 2 giây, Python/PyPy 5 giây.
8
1000 300 1500 700 1200 1400 800 400
1 2 900
1 3 700
3 4 700
3 5 900
1 6 800
5 7 700
6 8 5005Năm cặp ghép được: $\{1,3\}$ ($d = 700$), $\{1,6\}$ ($d = 800$), $\{3,4\}$ ($d = 700 = r_4$), $\{3,5\}$ ($d = 900$), $\{5,7\}$ ($d = 700$). Cặp $\{1,2\}$ không được vì $d = 900$ nhưng $r_2 = 300$; tương tự $\{3,6\}$ ($d = 1500 > r_6 = 1400$) và $\{6,8\}$ ($d = 500 > r_8 = 400$).
Bạn nào muốn thử nộp bài để hệ thống chấm xem có đúng hay không thì dùng link này nhé. Bộ test có chia subtask và có cả những test "bẫy", nên code đúng ý tưởng nhưng sai chi tiết nhỏ cũng sẽ bị trừ điểm đó.
Nộp bài Học nhómLời mở đầu
Bài này là một bài "đếm cặp đỉnh trên cây" – dạng bài kinh điển của kỹ thuật phân rã trọng tâm (centroid decomposition). Nhưng điều kiện ghép cặp có một "nút thắt" nhỏ mà rất khó chịu: nó không chỉ phụ thuộc vào khoảng cách $d(u, v)$, mà còn phụ thuộc vào cả hai bán kính $r_u$ và $r_v$ cùng lúc:
$$d(u, v) \le \min(r_u, r_v).$$Nếu mọi $r_i$ bằng nhau (subtask 3) thì đây đúng là bài kinh điển "đếm cặp đỉnh có khoảng cách không quá $R$". Khi các $r_i$ khác nhau, ta cần thêm một mẹo nữa. Hai nhận xét quan trọng của cả bài:
- Chỉ bạn có $r$ nhỏ hơn mới quyết định. Trong một cặp, điều kiện "$d \le r_u$ và $d \le r_v$" tương đương với "$d \le$ bán kính của bạn có $r$ nhỏ hơn". Nếu ta xét các bạn theo thứ tự $r$ giảm dần, thì khi tới bạn $u$, mọi bạn đã xét trước đều có $r \ge r_u$, nên với họ điều kiện chỉ còn là $d \le r_u$ – một điều kiện của riêng $u$!
- Mọi đường đi trên cây đều đi qua đúng một "trọng tâm" được chọn sớm nhất. Chọn trọng tâm $c$ của cây, đếm các cặp có đường đi qua $c$, xoá $c$ rồi làm tiếp trên từng mảnh. Với cặp $u, v$ ở hai nhánh khác nhau của $c$: $d(u, v) = d_u + d_v$, trong đó $d_x$ là khoảng cách từ $x$ tới $c$.
Trọng tâm $c$ chia cây thành các nhánh. Cặp khác nhánh có $d(u, v) = d_u + d_v$; cặp cùng nhánh thì $d_u + d_x$ chỉ là đường "vòng qua $c$" (dài hơn đường thật), sẽ được đếm đúng ở các tầng sâu hơn.
Mình sẽ đi qua 4 cách, mỗi cách ứng với một subtask, cách sau dùng lại ý của cách trước:
Cách 1DFS từ từng nhà – O(N²) · 20%
Cách tự nhiên nhất: với mỗi nhà $s$, chạy DFS (hoặc BFS) để tính khoảng cách $d(s, u)$ tới mọi nhà $u$, rồi kiểm tra trực tiếp $d(s, u) \le r_s$ và $d(s, u) \le r_u$. Để mỗi cặp chỉ đếm một lần, ta chỉ đếm khi $u > s$.
Mỗi lần DFS mất $O(N)$, làm $N$ lần nên tổng $O(N^2)$: với $N \le 2000$ là $4 \cdot 10^6$ phép tính, chạy trong nháy mắt. Với $N = 7 \cdot 10^4$ thì là $4{,}9 \cdot 10^9$ – quá chậm.
long long. Số cặp tối đa là $\frac{N(N-1)}{2} \approx 2{,}45 \cdot 10^9$, cũng vượt int!#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
int n; scanf("%d", &n);
vector<ll> r(n + 1);
for (int i = 1; i <= n; i++) scanf("%lld", &r[i]);
vector<vector<pair<int, ll>>> g(n + 1);
for (int i = 1; i < n; i++) {
int u, v; ll w; scanf("%d %d %lld", &u, &v, &w);
g[u].push_back({v, w}); g[v].push_back({u, w});
}
ll ans = 0;
vector<ll> d(n + 1); vector<int> par(n + 1), st;
for (int s = 1; s <= n; s++) { // DFS (bang ngan xep) tu nha s
d[s] = 0; par[s] = 0; st.assign(1, s);
while (!st.empty()) {
int u = st.back(); st.pop_back();
if (u > s && d[u] <= r[u] && d[u] <= r[s]) ans++; // moi cap dem 1 lan (u > s)
for (auto [v, w] : g[u]) if (v != par[u]) { par[v] = u; d[v] = d[u] + w; st.push_back(v); }
}
}
printf("%lld\n", ans);
return 0;
}import sys
def main():
data = sys.stdin.buffer.read().split()
n = int(data[0])
r = [0] * (n + 1)
for i in range(1, n + 1):
r[i] = int(data[i])
g = [[] for _ in range(n + 1)]
p = n + 1
for _ in range(n - 1):
u = int(data[p]); v = int(data[p + 1]); w = int(data[p + 2]); p += 3
g[u].append((v, w)); g[v].append((u, w))
ans = 0
d = [0] * (n + 1); par = [0] * (n + 1)
for s in range(1, n + 1): # DFS (bang ngan xep) tu nha s
d[s] = 0; par[s] = 0; st = [s]; rs = r[s]
while st:
u = st.pop(); du = d[u]
if u > s and du <= r[u] and du <= rs: # moi cap dem 1 lan (u > s)
ans += 1
pu = par[u]
for v, w in g[u]:
if v != pu:
par[v] = u; d[v] = du + w; st.append(v)
print(ans)
main()Cách 2Các nhà trên một đường thẳng: sắp theo r giảm dần + Fenwick – O(N log N) · 40%
Ở subtask 2, nhà $i$ nằm ở vị trí $x_i$ trên một đường thẳng ($x_1 = 0$, $x_{i+1} = x_i + w_i$), và $d(u, v) = |x_u - x_v|$. Điều kiện ghép cặp:
$$|x_u - x_v| \le r_u \quad\text{và}\quad |x_u - x_v| \le r_v.$$Khó ở chỗ có hai điều kiện, mỗi điều kiện gắn với một bạn. Mẹo quyết định của cả bài nằm ở đây:
Vậy thuật toán là: sắp xếp các bạn theo $r$ giảm dần; với mỗi bạn $u$ theo thứ tự đó:
- Tìm (bằng chặt nhị phân trên mảng $x$ tăng dần) đoạn chỉ số $[lo, hi]$ các nhà có vị trí trong $[x_u - r_u,\ x_u + r_u]$.
- Đếm số nhà đã được thêm có chỉ số trong $[lo, hi]$ – dùng cây Fenwick (BIT) trên chỉ số nhà – và cộng vào đáp án.
- Thêm nhà $u$ vào Fenwick.
Vì sao các bạn có $r$ bằng nhau không bị đếm sai? Giả sử $r_u = r_v$ và $v$ đứng trước $u$ trong thứ tự. Khi xét $u$, ta kiểm tra $d \le r_u = \min(r_u, r_v)$ – đúng điều kiện. Còn khi xét $v$ thì $u$ chưa được thêm nên không bị đếm lần hai. Điều quan trọng là: hỏi trước, thêm sau, từng bạn một.
Các em bấm thử mô phỏng dưới đây. Mỗi bước xét một bạn (theo $r$ giảm dần): ô vàng là bạn đang xét, ô xanh lá là các bạn đã thêm nằm trong đoạn với tới (được đếm), ô xanh dương là các bạn đã thêm nhưng ở quá xa.
Trong code, với các nhà không nằm trên đường thẳng (subtask 1) ta vẫn dùng lại Cách 1. Mảng $x$ tăng ngặt (vì $w \ge 1$), nên lower_bound/upper_bound trên $x$ cho ngay đoạn chỉ số $[lo, hi]$.
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n;
vector<ll> r;
vector<vector<pair<int, ll>>> g;
ll brute() { // Cach 1: DFS tu tung nha, O(N^2)
ll ans = 0;
vector<ll> d(n + 1); vector<int> par(n + 1), st;
for (int s = 1; s <= n; s++) {
d[s] = 0; par[s] = 0; st.assign(1, s);
while (!st.empty()) {
int u = st.back(); st.pop_back();
if (u > s && d[u] <= r[u] && d[u] <= r[s]) ans++;
for (auto [v, w] : g[u]) if (v != par[u]) { par[v] = u; d[v] = d[u] + w; st.push_back(v); }
}
}
return ans;
}
ll path() { // duong thang: nha i o vi tri x[i]
vector<ll> x(n + 1, 0);
for (int i = 1; i < n; i++) for (auto [v, w] : g[i]) if (v == i + 1) x[i + 1] = x[i] + w;
vector<int> id(n); iota(id.begin(), id.end(), 1);
sort(id.begin(), id.end(), [&](int a, int b) { return r[a] > r[b]; });
vector<int> bit(n + 1, 0);
ll res = 0;
for (int u : id) { // cac ban da them deu co r >= r[u]: chi can |x[v] - x[u]| <= r[u]
int lo = lower_bound(x.begin() + 1, x.end(), x[u] - r[u]) - x.begin();
int hi = upper_bound(x.begin() + 1, x.end(), x[u] + r[u]) - x.begin() - 1;
for (int k = hi; k > 0; k -= k & -k) res += bit[k];
for (int k = lo - 1; k > 0; k -= k & -k) res -= bit[k];
for (int k = u; k <= n; k += k & -k) bit[k]++;
}
return res;
}
int main() {
scanf("%d", &n);
r.assign(n + 1, 0); g.assign(n + 1, {});
for (int i = 1; i <= n; i++) scanf("%lld", &r[i]);
bool isPath = true;
for (int i = 1; i < n; i++) {
int u, v; ll w; scanf("%d %d %lld", &u, &v, &w);
g[u].push_back({v, w}); g[v].push_back({u, w});
if (abs(u - v) != 1) isPath = false;
}
printf("%lld\n", isPath ? path() : brute());
return 0;
}import sys
from bisect import bisect_left, bisect_right
def brute(n, r, g): # Cach 1: DFS tu tung nha, O(N^2)
ans = 0
d = [0] * (n + 1); par = [0] * (n + 1)
for s in range(1, n + 1):
d[s] = 0; par[s] = 0; st = [s]; rs = r[s]
while st:
u = st.pop(); du = d[u]
if u > s and du <= r[u] and du <= rs:
ans += 1
pu = par[u]
for v, w in g[u]:
if v != pu:
par[v] = u; d[v] = du + w; st.append(v)
return ans
def path(n, r, g): # duong thang: nha i o vi tri x[i]
x = [0] * (n + 1)
for i in range(1, n):
for v, w in g[i]:
if v == i + 1:
x[i + 1] = x[i] + w
ids = sorted(range(1, n + 1), key=r.__getitem__, reverse=True)
bit = [0] * (n + 1)
res = 0
for u in ids: # cac ban da them deu co r >= r[u]: chi can |x[v] - x[u]| <= r[u]
lo = bisect_left(x, x[u] - r[u], 1)
k = bisect_right(x, x[u] + r[u], 1) - 1
while k:
res += bit[k]; k &= k - 1
k = lo - 1
while k:
res -= bit[k]; k &= k - 1
k = u
while k <= n:
bit[k] += 1; k += k & -k
return res
def main():
data = sys.stdin.buffer.read().split()
n = int(data[0])
r = [0] * (n + 1)
for i in range(1, n + 1):
r[i] = int(data[i])
g = [[] for _ in range(n + 1)]
is_path = True
p = n + 1
for _ in range(n - 1):
u = int(data[p]); v = int(data[p + 1]); w = int(data[p + 2]); p += 3
g[u].append((v, w)); g[v].append((u, w))
if abs(u - v) != 1:
is_path = False
print(path(n, r, g) if is_path else brute(n, r, g))
main()Cách 3Mọi r bằng nhau: phân rã trọng tâm + hai con trỏ – O(N log² N) · 65%
Khi $r_1 = \ldots = r_N = R$, bài toán trở thành: đếm số cặp đỉnh của cây có khoảng cách không quá $R$. Đây là bài mẫu của phân rã trọng tâm.
Trọng tâm của một cây $m$ đỉnh là đỉnh $c$ mà khi xoá đi, mọi mảnh còn lại có không quá $m / 2$ đỉnh. Cây nào cũng có trọng tâm: tính kích thước cây con $sz$ bằng một lần duyệt, rồi tìm đỉnh có "mảnh lớn nhất" $\max(m - sz_u, \max_{con\ v} sz_v) \le m / 2$.
Phân rã: với cây (thành phần) hiện tại:
- Tìm trọng tâm $c$; tính $d_x$ = khoảng cách từ $c$ tới mọi đỉnh $x$ trong thành phần.
- Đếm các cặp $\{u, v\}$ trong thành phần có $d_u + d_v \le R$ (kể cả cặp có một đầu là $c$, vì $d_c = 0$).
- Cặp nằm cùng một nhánh của $c$ thì $d_u + d_v$ không phải độ dài thật (xem hình ở trên). Với từng nhánh, đếm các cặp trong nhánh đó có $d_u + d_v \le R$ (vẫn dùng $d$ tính từ $c$) rồi trừ đi. Sau bước này, ta đã đếm đúng mọi cặp có đường đi qua $c$.
- Xoá $c$, làm tiếp đệ quy với từng nhánh (các cặp còn lại nằm gọn trong một nhánh).
Vì mỗi lần xuống một tầng, kích thước thành phần giảm ít nhất một nửa, nên có không quá $\log_2 N$ tầng, và mỗi đỉnh xuất hiện trong $O(\log N)$ thành phần. Bước đếm "cặp có tổng $\le R$" làm bằng sắp xếp + hai con trỏ: sắp $d$ tăng dần, con trỏ $i$ từ trái, $j$ từ phải; với mỗi $i$, lùi $j$ tới khi $d_i + d_j \le R$, rồi cộng $j - i$. Tổng cộng $O(N \log^2 N)$.
Code dưới đây giữ lại Cách 1 và Cách 2 cho các subtask trước, và dùng phân rã trọng tâm khi mọi $r$ bằng nhau. Toàn bộ được viết không đệ quy (dùng ngăn xếp và BFS) để không lo tràn ngăn xếp khi cây là một đường dài.
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n;
vector<ll> r;
vector<vector<pair<int, ll>>> g;
vector<ll> dep;
vector<int> par, sz, ord;
vector<char> del;
ll brute() { // Cach 1: DFS tu tung nha, O(N^2)
ll ans = 0;
vector<ll> d(n + 1); vector<int> par(n + 1), st;
for (int s = 1; s <= n; s++) {
d[s] = 0; par[s] = 0; st.assign(1, s);
while (!st.empty()) {
int u = st.back(); st.pop_back();
if (u > s && d[u] <= r[u] && d[u] <= r[s]) ans++;
for (auto [v, w] : g[u]) if (v != par[u]) { par[v] = u; d[v] = d[u] + w; st.push_back(v); }
}
}
return ans;
}
ll path() { // duong thang: nha i o vi tri x[i]
vector<ll> x(n + 1, 0);
for (int i = 1; i < n; i++) for (auto [v, w] : g[i]) if (v == i + 1) x[i + 1] = x[i] + w;
vector<int> id(n); iota(id.begin(), id.end(), 1);
sort(id.begin(), id.end(), [&](int a, int b) { return r[a] > r[b]; });
vector<int> bit(n + 1, 0);
ll res = 0;
for (int u : id) { // cac ban da them deu co r >= r[u]: chi can |x[v] - x[u]| <= r[u]
int lo = lower_bound(x.begin() + 1, x.end(), x[u] - r[u]) - x.begin();
int hi = upper_bound(x.begin() + 1, x.end(), x[u] + r[u]) - x.begin() - 1;
for (int k = hi; k > 0; k -= k & -k) res += bit[k];
for (int k = lo - 1; k > 0; k -= k & -k) res -= bit[k];
for (int k = u; k <= n; k += k & -k) bit[k]++;
}
return res;
}
void bfs(int s, int p, ll d0) { // BFS tren cac dinh chua bi xoa
ord.assign(1, s); par[s] = p; dep[s] = d0;
for (int i = 0; i < (int)ord.size(); i++) {
int u = ord[i];
for (auto [v, w] : g[u])
if (v != par[u] && !del[v]) { par[v] = u; dep[v] = dep[u] + w; ord.push_back(v); }
}
}
ll calcEq(ll R) { // so cap trong ord co dep[u] + dep[v] <= R
vector<ll> ds;
for (int u : ord) ds.push_back(dep[u]);
sort(ds.begin(), ds.end());
ll res = 0;
int j = ds.size() - 1;
for (int i = 0; i < j; i++) { // hai con tro
while (j > i && ds[i] + ds[j] > R) j--;
res += j - i;
}
return res;
}
ll centroid() { // Cach 3: moi r bang nhau
dep.assign(n + 1, 0); par.assign(n + 1, 0); sz.assign(n + 1, 0); del.assign(n + 1, 0);
ll R = r[1], ans = 0;
vector<int> st = {1};
while (!st.empty()) {
int s = st.back(); st.pop_back();
bfs(s, 0, 0);
int tot = ord.size(), c = s;
for (int i = tot - 1; i >= 0; i--) {
int u = ord[i]; sz[u] = 1;
for (auto [v, w] : g[u]) if (v != par[u] && !del[v]) sz[u] += sz[v];
}
for (int u : ord) {
int mx = tot - sz[u];
for (auto [v, w] : g[u]) if (v != par[u] && !del[v]) mx = max(mx, sz[v]);
if (2 * mx <= tot) { c = u; break; }
}
bfs(c, 0, 0);
ans += calcEq(R);
del[c] = 1;
for (auto [v, w] : g[c]) if (!del[v]) {
bfs(v, c, w);
ans -= calcEq(R);
st.push_back(v);
}
}
return ans;
}
int main() {
scanf("%d", &n);
r.assign(n + 1, 0); g.assign(n + 1, {});
for (int i = 1; i <= n; i++) scanf("%lld", &r[i]);
bool isPath = true, same = true;
for (int i = 1; i <= n; i++) if (r[i] != r[1]) same = false;
for (int i = 1; i < n; i++) {
int u, v; ll w; scanf("%d %d %lld", &u, &v, &w);
g[u].push_back({v, w}); g[v].push_back({u, w});
if (abs(u - v) != 1) isPath = false;
}
printf("%lld\n", same ? centroid() : isPath ? path() : brute());
return 0;
}import sys
from bisect import bisect_left, bisect_right
def brute(n, r, g): # Cach 1: DFS tu tung nha, O(N^2)
ans = 0
d = [0] * (n + 1); par = [0] * (n + 1)
for s in range(1, n + 1):
d[s] = 0; par[s] = 0; st = [s]; rs = r[s]
while st:
u = st.pop(); du = d[u]
if u > s and du <= r[u] and du <= rs:
ans += 1
pu = par[u]
for v, w in g[u]:
if v != pu:
par[v] = u; d[v] = du + w; st.append(v)
return ans
def path(n, r, g): # duong thang: nha i o vi tri x[i]
x = [0] * (n + 1)
for i in range(1, n):
for v, w in g[i]:
if v == i + 1:
x[i + 1] = x[i] + w
ids = sorted(range(1, n + 1), key=r.__getitem__, reverse=True)
bit = [0] * (n + 1)
res = 0
for u in ids: # cac ban da them deu co r >= r[u]: chi can |x[v] - x[u]| <= r[u]
lo = bisect_left(x, x[u] - r[u], 1)
k = bisect_right(x, x[u] + r[u], 1) - 1
while k:
res += bit[k]; k &= k - 1
k = lo - 1
while k:
res -= bit[k]; k &= k - 1
k = u
while k <= n:
bit[k] += 1; k += k & -k
return res
def centroid(n, r, g): # Cach 3: moi r bang nhau
R = r[1]
dep = [0] * (n + 1); par = [0] * (n + 1); sz = [0] * (n + 1)
dead = [False] * (n + 1)
def bfs(s, pa, d0): # BFS tren cac dinh chua bi xoa
order = [s]; par[s] = pa; dep[s] = d0
i = 0
while i < len(order):
u = order[i]; i += 1
pu = par[u]; du = dep[u]
for v, w in g[u]:
if v != pu and not dead[v]:
par[v] = u; dep[v] = du + w; order.append(v)
return order
def calc_eq(order): # so cap co dep[u] + dep[v] <= R (hai con tro)
ds = sorted([dep[u] for u in order])
res = 0; j = len(ds) - 1; i = 0
while i < j:
while j > i and ds[i] + ds[j] > R:
j -= 1
res += j - i; i += 1
return res
ans = 0
st = [1]
while st:
s = st.pop()
order = bfs(s, 0, 0)
tot = len(order); c = s
for i in range(tot - 1, -1, -1):
u = order[i]; t = 1; pu = par[u]
for v, w in g[u]:
if v != pu and not dead[v]:
t += sz[v]
sz[u] = t
for u in order:
mx = tot - sz[u]; pu = par[u]
for v, w in g[u]:
if v != pu and not dead[v] and sz[v] > mx:
mx = sz[v]
if 2 * mx <= tot:
c = u; break
ans += calc_eq(bfs(c, 0, 0))
dead[c] = True
for v, w in g[c]:
if not dead[v]:
ans -= calc_eq(bfs(v, c, w))
st.append(v)
return ans
def main():
data = sys.stdin.buffer.read().split()
n = int(data[0])
r = [0] * (n + 1)
for i in range(1, n + 1):
r[i] = int(data[i])
g = [[] for _ in range(n + 1)]
is_path = True
same = all(r[i] == r[1] for i in range(1, n + 1))
p = n + 1
for _ in range(n - 1):
u = int(data[p]); v = int(data[p + 1]); w = int(data[p + 2]); p += 3
g[u].append((v, w)); g[v].append((u, w))
if abs(u - v) != 1:
is_path = False
print(centroid(n, r, g) if same else path(n, r, g) if is_path else brute(n, r, g))
main()Cách 4Phân rã trọng tâm + sắp theo r + Fenwick – O(N log² N) · 100%
Bây giờ ghép hai ý lại. Khung phân rã trọng tâm giữ nguyên như Cách 3; chỉ thay bước "đếm cặp trong một danh sách". Với trọng tâm $c$ và một danh sách đỉnh (cả thành phần, hoặc một nhánh), mỗi đỉnh $x$ có hai số $(r_x, d_x)$, ta cần đếm số cặp
$$d_u + d_v \le \min(r_u, r_v).$$Áp dụng mẹo của Cách 2: sắp các đỉnh theo $r$ giảm dần, xét lần lượt. Khi tới $u$, các đỉnh $v$ đã thêm đều có $r_v \ge r_u$, nên điều kiện thành
$$d_v \le r_u - d_u.$$Tức là: đếm số đỉnh đã thêm có $d_v$ không vượt quá ngưỡng $r_u - d_u$. Ngưỡng này không đơn điệu nên không dùng hai con trỏ được nữa – ta dùng cây Fenwick trên giá trị $d$ đã nén (sắp $d$ tăng dần thành mảng ds, vị trí của một giá trị tìm bằng chặt nhị phân):
k = upper_bound(ds, r_u - d_u): số giá trị $d$ (trong danh sách) không vượt ngưỡng; tổng tiền tố Fenwick tới $k$ là số đỉnh đã thêm thoả điều kiện.- Thêm $u$: tăng Fenwick tại vị trí
lower_bound(ds, d_u) + 1.
Phần trừ các cặp cùng nhánh làm y hệt Cách 3: gọi đúng hàm đếm đó trên từng nhánh (với $d$ vẫn tính từ $c$) rồi trừ đi. Lý do vẫn đúng: một cặp cùng nhánh được hàm đếm trên cả thành phần tính vào khi và chỉ khi nó được hàm đếm trên nhánh tính vào (cùng $r$, cùng $d$), nên hai lần đếm triệt tiêu nhau.
- Hỏi trước, thêm sau – từng đỉnh một. Nếu thêm cả nhóm các đỉnh có cùng $r$ rồi mới hỏi, các cặp cùng $r$ sẽ bị đếm hai lần (subtask 3 toàn là $r$ bằng nhau!).
- Đừng quên phần trừ các cặp cùng nhánh.
- long long cho $d$, $r - d$ và đáp án.
Mô phỏng dưới đây chạy toàn bộ quá trình phân rã trên một cây nhỏ: mỗi trọng tâm có một bước "cộng" (đếm trên cả thành phần) và các bước "trừ" (đếm trên từng nhánh). Trong bảng, các đỉnh đã được sắp theo $r$ giảm dần; cột cuối là số đỉnh phía trên (đã thêm) có $d \le r - d$ của dòng đó.
Với ví dụ đề: trọng tâm đầu tiên là nhà $1$, đếm được $2$ cặp ($\{1, 3\}$ và $\{1, 6\}$); tiếp theo trong nhánh $\{3, 4, 5, 7\}$, trọng tâm nhà $3$ đếm thêm $\{3, 4\}$, $\{3, 5\}$; cuối cùng trọng tâm nhà $5$ đếm $\{5, 7\}$. Tổng là $5$. Ở ví dụ đề phần "trừ" luôn bằng $0$ – các em bấm ví dụ "Đường thẳng, r bằng nhau" để thấy phần trừ hoạt động: với trọng tâm nhà $3$, cặp $\{4, 5\}$ có $d_4 + d_5 = 2 + 4 = 6 \le 6$ nên bị đếm "vòng qua nhà $3$" và phải trừ đi (cặp này được đếm lại đúng ở tầng sau).
Code hoàn chỉnh (hàm calc là phần mới so với Cách 3). Bản Python cùng ý tưởng nhưng được tối ưu để chạy kịp trên PyPy: mỗi trọng tâm chỉ BFS một lần và gom luôn các nhánh, tìm trọng tâm bằng cách "đi xuống" nhánh con quá nửa, xếp hạng theo $r$ một lần cho cả bài, và với danh sách nhỏ (≤ 20 000 đỉnh) thì dùng mảng sắp xếp + insort thay cho Fenwick.
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n;
vector<vector<pair<int, ll>>> g;
vector<ll> r, dep;
vector<int> par, sz, ord;
vector<char> del;
// BFS tu s (cha p, khoang cach ban dau d0) tren cac dinh chua bi xoa
void bfs(int s, int p, ll d0) {
ord.assign(1, s); par[s] = p; dep[s] = d0;
for (int i = 0; i < (int)ord.size(); i++) {
int u = ord[i];
for (auto [v, w] : g[u])
if (v != par[u] && !del[v]) { par[v] = u; dep[v] = dep[u] + w; ord.push_back(v); }
}
}
// dem so cap trong ord co dep[u] + dep[v] <= min(r[u], r[v])
ll calc() {
int m = ord.size();
vector<ll> ds(m);
vector<pair<ll, ll>> a(m); // (r, dep) cua tung ban
for (int i = 0; i < m; i++) { ds[i] = dep[ord[i]]; a[i] = {r[ord[i]], dep[ord[i]]}; }
sort(ds.begin(), ds.end());
sort(a.rbegin(), a.rend()); // r giam dan
vector<int> bit(m + 1, 0);
ll res = 0;
for (auto [ru, du] : a) { // cac ban da them deu co r >= ru: chi can dep[v] <= ru - du
int k = upper_bound(ds.begin(), ds.end(), ru - du) - ds.begin();
for (; k > 0; k -= k & -k) res += bit[k];
int p = lower_bound(ds.begin(), ds.end(), du) - ds.begin() + 1;
for (; p <= m; p += p & -p) bit[p]++;
}
return res;
}
int main() {
scanf("%d", &n);
g.assign(n + 1, {}); r.assign(n + 1, 0); dep.assign(n + 1, 0);
par.assign(n + 1, 0); sz.assign(n + 1, 0); del.assign(n + 1, 0);
for (int i = 1; i <= n; i++) scanf("%lld", &r[i]);
for (int i = 1; i < n; i++) {
int u, v; ll w; scanf("%d %d %lld", &u, &v, &w);
g[u].push_back({v, w}); g[v].push_back({u, w});
}
ll ans = 0;
vector<int> st = {1};
while (!st.empty()) {
int s = st.back(); st.pop_back();
bfs(s, 0, 0);
int tot = ord.size(), c = s;
for (int i = tot - 1; i >= 0; i--) {
int u = ord[i]; sz[u] = 1;
for (auto [v, w] : g[u]) if (v != par[u] && !del[v]) sz[u] += sz[v];
}
for (int u : ord) { // tim trong tam: moi phan con lai sau khi xoa u co <= tot/2 dinh
int mx = tot - sz[u];
for (auto [v, w] : g[u]) if (v != par[u] && !del[v]) mx = max(mx, sz[v]);
if (2 * mx <= tot) { c = u; break; }
}
bfs(c, 0, 0);
ans += calc(); // moi cap trong thanh phan, tinh duong vong qua c
del[c] = 1;
for (auto [v, w] : g[c]) if (!del[v]) {
bfs(v, c, w);
ans -= calc(); // bo cap nam cung mot nhanh (duong that khong qua c)
st.push_back(v);
}
}
printf("%lld\n", ans);
return 0;
}import sys
from bisect import bisect_right, bisect_left, insort
def main():
data = sys.stdin.buffer.read().split()
n = int(data[0])
r = [0] * (n + 1)
for i in range(1, n + 1):
r[i] = int(data[i])
adj = [[] for _ in range(n + 1)]
wt = [[] for _ in range(n + 1)]
p = n + 1
for _ in range(n - 1):
u = int(data[p]); v = int(data[p + 1]); w = int(data[p + 2]); p += 3
adj[u].append(v); wt[u].append(w)
adj[v].append(u); wt[v].append(w)
# xep hang theo r giam dan (dung chung cho moi lan dem)
byr = sorted(range(1, n + 1), key=r.__getitem__, reverse=True)
rk = [0] * (n + 1)
for i in range(n):
rk[byr[i]] = i
dep = [0] * (n + 1)
par = [0] * (n + 1)
sz = [0] * (n + 1)
br = [0] * (n + 1)
dead = [False] * (n + 1)
def calc(lst):
# dem so cap trong lst co dep[u] + dep[v] <= min(r[u], r[v])
m = len(lst)
if m < 2:
return 0
ks = sorted([rk[u] for u in lst])
res = 0
if m <= 20000:
L = []
for k in ks:
u = byr[k]
du = dep[u]
t = r[u] - du
if t >= 0:
res += bisect_right(L, t)
insort(L, du)
return res
ds = sorted([dep[u] for u in lst])
bit = [0] * (m + 1)
for k in ks: # cac ban da them deu co r >= r[u]
u = byr[k]
du = dep[u]
t = r[u] - du
if t >= 0:
j = bisect_right(ds, t)
while j:
res += bit[j]; j &= j - 1
q = bisect_left(ds, du) + 1
while q <= m:
bit[q] += 1; q += q & -q
return res
# thanh phan ban dau: BFS tu dinh 1
order = [1]; par[1] = 0
i = 0
while i < len(order):
u = order[i]; i += 1
pu = par[u]
for v in adj[u]:
if v != pu:
par[v] = u; order.append(v)
ans = 0
st = [order]
while st:
order = st.pop()
tot = len(order)
if tot == 1:
dead[order[0]] = True
continue
# kich thuoc cay con theo thu tu BFS (cha dung truoc con)
for u in order:
sz[u] = 1
for idx in range(tot - 1, 0, -1):
u = order[idx]
sz[par[u]] += sz[u]
# tim trong tam: di xuong nhanh con co kich thuoc > tot/2
c = order[0]; half = tot >> 1
while True:
pc = par[c]; nx = 0
for v in adj[c]:
if v != pc and not dead[v] and sz[v] > half:
nx = v; break
if nx == 0:
break
c = nx
# mot lan BFS tu trong tam, gan nhan nhanh
dead[c] = True
par[c] = 0; dep[c] = 0
groups = []
o2 = [c]
ac = adj[c]; wc = wt[c]
for j in range(len(ac)):
v = ac[j]
if not dead[v]:
par[v] = c; dep[v] = wc[j]; br[v] = len(groups)
groups.append([v]); o2.append(v)
i = 1
while i < len(o2):
u = o2[i]; i += 1
pu = par[u]; du = dep[u]; wu = wt[u]; au = adj[u]; b = br[u]
g = groups[b]
for j in range(len(au)):
v = au[j]
if v != pu and not dead[v]:
par[v] = u; dep[v] = du + wu[j]; br[v] = b
o2.append(v); g.append(v)
ans += calc(o2)
for g in groups:
ans -= calc(g)
par[g[0]] = 0
st.append(g)
sys.stdout.write(str(ans) + "\n")
main()Các em chạy tay từng dòng code C++ ở trên với dữ liệu nhỏ: bảng bên phải hiện mảng dep, các đỉnh đã xoá, danh sách ord của lần BFS gần nhất và bảng Fenwick trong hàm calc.
Bộ test chấm "bắt" những lỗi nào?
| Cách làm | Sub 1 (20%, 10 test) | Sub 2 (20%, 7 test) | Sub 3 (25%, 8 test) | Sub 4 (35%, 11 test) | Điểm |
|---|---|---|---|---|---|
| Cách 1 – DFS từ từng nhà | Đúng | TLE | TLE | TLE | 20 |
| Cách 2 – đường thẳng: sắp theo r + Fenwick | Đúng | Đúng | TLE | TLE | 40 |
| Cách 3 – r bằng nhau: centroid + hai con trỏ | Đúng | Đúng | Đúng | TLE | 65 |
| Cách 4 – centroid + sắp theo r + Fenwick (lời giải chuẩn) | Đúng | Đúng | Đúng | Đúng | 100 |
Sai: khoảng cách và biến đếm kiểu int | 6/10 | 3/7 | 3/8 | 1/11 | 33,13 |
| Sai: tách tại đỉnh bất kỳ thay vì trọng tâm | Đúng | TLE | 2/8 | 1/11 | 29,43 |
| Sai: thêm cả nhóm r bằng nhau rồi mới đếm (đếm trùng) | 4/10 | 1/7 | 0/8 | 2/11 | 17,22 |
| Sai: quên trừ các cặp cùng nhánh | 3/10 | 0/7 | 0/8 | 0/11 | 6 |
Ghi chú của mình: các cách C++ ở trên đã được chấm thử với đúng bộ test của bài: bộ test có 36 test (10 + 7 + 8 + 11), mỗi cách đạt đúng số điểm công bố, còn bốn lời giải sai bị bắt như bảng trên. Lời giải chuẩn sol.cpp chạy test lâu nhất khoảng 0,6 giây (giới hạn 2 giây); bản Python sol.py cho cùng đáp án trên mọi test (CPython khoảng 4 giây, các em nhớ chọn PyPy 3 khi nộp). "TLE" nghĩa là chạy quá 3 lần giới hạn thời gian. Để ý hai lỗi "đau" nhất: quên long long mất gần hết điểm vì khoảng cách lên tới cỡ $7 \cdot 10^{13}$ và số cặp lên tới cỡ $2{,}4 \cdot 10^9$; còn tách cây tại một đỉnh bất kỳ (không phải trọng tâm) thì vẫn đúng đáp án nhưng trên các cây dài (đường thẳng, cây "sâu róm", cây "con nhện") sẽ chạy quá chậm – chỉ còn qua được vài test cây "bè" như hình sao.
Rất vui khi anh em nào đã đọc đến những dòng cuối cùng này. Anh em đọc chắc cũng mỏi mắt luôn chứ 😅 vậy anh em nghĩ người viết ra hết đống ở trên còn thế nào nữa 🤪. Vậy nên đừng tiếc gì cho mình một lời bình luận hoặc chia sẻ bài viết để động viên mình nè. Chúc anh em học tốt và gặt hái nhiều thành công 😇
- Không biết học gì trước, học gì sau, học lan man nên kiến thức bị hổng.
- Bí bài không ai gỡ, ngồi cả buổi vẫn không ra rồi bỏ cuộc.
- Code chạy được nhưng không đủ điểm, không biết sai ở đâu, chậm ở đâu.
- Không có đề thi thử và người chấm, nên không biết mình đang ở mức nào.
- Dễ nản giữa chừng khi không có ai theo sát.
- Lộ trình 111 buổi video sắp sẵn từ dòng code đầu tiên đến đề HSG thật của các tỉnh.
- Thầy trực tiếp gỡ bí, chữa bài cho từng em qua nhóm Zalo riêng.
- Hệ thống chấm bài tự động 24/7 và thi thử HSG hằng tuần để biết chính xác mình đang ở đâu.
- Đóng học phí một lần, thầy kèm lâu dài: từ thi chuyên Tin đến HSG quốc gia và cả khi lên đại học.
Nhận xét
Đăng nhận xét