---
# System prepended metadata

title: Tìm kho báu

---

# Tìm kho báu

## Subtask 1

Với $k=1$, bài toán trở về tìm đường đi dài nhất trên cây. Ta có thể $dfs$ $2$ lần để giải quyết vấn đề này.

## Subtask 2

Với $n \le 20$, ta có thể quay lui duyệt thử toàn bộ các trường hợp có thể đi. 

## Subtask 3

Với $n \le 300$ và $k = 2$, ta sẽ cần duyệt vị trí xuất phát và coi như đó là gốc của cây, sau đó áp dụng quy hoạch động trên cây với đỉnh xuất phát đang xét như sau:
- Gọi $dp[u][0]$ là đường đi dài nhất từ $u$ đi xuống các đỉnh con của $u$ và **không quay lại đỉnh cha** của $u$.
- Gọi $dp[u][1]$ là đường đi dài nhất từ $u$ đi xuống các đỉnh con của $u$ và **quay ngược lại đỉnh cha** của $u$.
- Vì $k=2$, nên nếu từ $u$ di chuyển rồi quay lại cha thì $u$ chỉ có thể di chuyển xuống một đỉnh con duy nhất do đó $dp[u][1] = max (dp[v][1]) + 2$ với $v$ là đỉnh con trực tiếp của $u$. Còn với trường hợp $u$ không quay lại cha thì ta sẽ chọn một đỉnh $v_1$ để đi xuống rồi quay trở lại, sau đó đi xuống $v_2$ và không quay trở lại nữa, do đó $dp[u][0] = max(dp[v_1][1] + dp[v_2][0]) + 3$.
- Vì phải duyệt $v_1$ và $v_2$ do đó độ phức tạp của mỗi lần cố định đỉnh xuất phát sẽ là $O(n^2)$, kéo theo độ phức tạp tổng là $O(n^3)$.

## Subtask 4

Với $n \le 1000$, ta cũng có trạng thái quy hoạch động như trên, nhưng vì $k$ có thể lớn hơn $2$ nên ta có:
- $dp[u][1] = max(dp[v_1][1] + dp[v_2][1] + ... + dp[v_{k-1}][1]) + 2\times(k-1)$. Ta sẽ cần chọn ra $v_1,...,v_{k-1}$ sao cho biểu thức trên là lớn nhất, nói cách khác ta chỉ cần chọn ra $k-1$ giá trị lớn nhất là xong.
- $dp[u][0] = max(dp[v_1][1] + dp[v_2][1] + ... + dp[v_{k-1}][1] + dp[v_k][0]) + 2\times (k-1)$. Để tính được  biểu thức này, ta sẽ sắp xếp theo $dp[v][1]$ và dùng mảng cộng dồn để duyệt lần lượt $v_k$ và tính được tổng lớn nhất của $k-1$ đỉnh còn lại.
- Như vậy ta thu được độ thuật toán có độ phức tạp $O(n^2 \times log_n)$ 

## Subtask 5

Tương tự subtask $4$ nhưng ta dùng quy hoạch động đảo gốc để giảm độ phức tạp tính toán.

## Code mẫu: (Subtask $1$, $2$, $4$):

```cpp=
#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1e5+5;
vector<int> a[N];
int s[N];
#define ii pair<int,int>
vector<int> b[N];
int n,k;
namespace sub1 {
	int lv[N];
	void dfs (int u, int p) {
		for (int i=0; i<a[u].size(); i++){
			int v = a[u][i];
			if (v == p) continue;
			lv[v] = lv[u] + 1;
			dfs (v,u);
		}
	}
	void solve () {
		dfs(1,1);
		int mx = 1;
		for (int i=1; i<=n; i++) if (lv[mx] < lv[i]) mx = i;
		lv[mx] = 0;
		dfs (mx,mx);
		int res = 0;
		for (int i=1; i<=n; i++) res = max (res,lv[i]);
		cout << res << '\n';
	}
}
namespace sub2 {
	int ans = 0;
	int dak[N],mim[N];
	void go (int u, int id, int res) {
		dak[u]--;
		mim[id]--;
		for (int i=0; i<a[u].size(); i++) {
			int v = a[u][i];
			int c = b[u][i];
			if (dak[v] > 0 && mim[c] > 0) go(v,c,res+1);
		}
		dak[u]++;
		mim[id]++;
		ans = max (ans,res);
	}
	void solve () {
		for (int i=1; i<=n; i++) dak[i] = k, mim[i-1] = 2;
		for (int i=1; i<=n; i++) go (i,0,0);
		cout << ans << '\n';
	}
}
namespace sub4 {
	int f[N][2];
	void dfs (int u, int p) {
		vector<ii> cur;
		f[u][0] = f[u][1] = 0;
		for (int i=0; i<a[u].size(); i++) {
			int v = a[u][i];
			if (v == p) continue;
			dfs (v,u);
			cur.push_back({f[v][1],f[v][0]});
		}
		sort (cur.begin(),cur.end(),greater<ii>());
		int g = cur.size();
		int add = min (g,k);
		int x = min (g,k-1);
		for (int i=0; i<cur.size(); i++) {
			s[i+1] = s[i] + cur[i].first;
		}
		f[u][1] = s[x] + 2*x;
		for (int i=1; i<=g; i++) {
			int val=0;
			if (i<k) {
				val = s[i-1] + s[add] - s[i] + cur[i-1].second + 2*add-1;
			}
			else val = s[k-1] + cur[i-1].second + 2*k-1;
			f[u][0] = max (f[u][0],val);
		}
	}
	void solve () {
		int res = 0;
		for (int i=1; i<=n; i++) {
			dfs (i,i);
			res = max (res,max (f[i][0],f[i][1]));
		}
		cout << res << '\n';
	}
}
signed main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	cin >> n >> k;
	for (int i=1; i<n; i++) {
		int u,v; cin >> u >> v;
		a[u].push_back(v);
		a[v].push_back(u);
		b[u].push_back(i);
		b[v].push_back(i);
	}
	if (k == 1) sub1::solve();
	else if (n<=18) sub2::solve();
	else sub4::solve(); 
}
```


## Code mẫu subtask $5$

```cpp=
#include<bits/stdc++.h>
using namespace std;
#define int long long
const int MAXN = 1e5 + 6.9;
int n, k;
vector<int> adj[MAXN];
int dp[MAXN][2];
int ans;
struct cmp0 {
    bool operator () (const int& x, const int& y) {
        if (dp[x][0] == dp[y][0]) return x < y;
        return dp[x][0] < dp[y][0];
    }
};
struct cmp1 {
    bool operator () (const int& x, const int& y) {
        if (dp[x][1] == dp[y][1]) return x < y;
        return dp[x][1] < dp[y][1];
    }
};
struct cmp01 {
    bool operator () (const int& x, const int& y) {
        if (dp[x][0] - dp[x][1] == dp[y][0] - dp[y][1]) return x < y;
        return dp[x][0] - dp[x][1] < dp[y][0] - dp[y][1];
    }
};
set<int, cmp0> s0[MAXN];
set<int, cmp1> t1[MAXN], s1[MAXN];
set<int, cmp01> t01[MAXN];
int sum[MAXN];
void Add(int u, int v) {
    t1[u].insert(v);
    t01[u].insert(v);
    sum[u] += dp[v][1];
    if (t1[u].size() >= k) {
        int tmp = *t1[u].begin();
        t1[u].erase(tmp);
        t01[u].erase(tmp);
        s1[u].insert(tmp);
        s0[u].insert(tmp);
        sum[u] -= dp[tmp][1];
    }
}
void Remove(int u, int v) {
    if (t1[u].find(v) != t1[u].end()) {
        t1[u].erase(v);
        t01[u].erase(v);
        sum[u] -= dp[v][1];
        if (!s1[u].empty()) {
            int tmp = *s1[u].rbegin();
            s1[u].erase(tmp);
            s0[u].erase(tmp);
            t1[u].insert(tmp);
            t01[u].insert(tmp);
            sum[u] += dp[tmp][1];
        }
    }
    else {
        s0[u].erase(v);
        s1[u].erase(v);
    }
}
void Calc(int u) {
    dp[u][1] = sum[u] + t1[u].size() * 2;
    int try1 = dp[u][1], try2 = dp[u][1];
    if (!s0[u].empty()) try1 += dp[*s0[u].rbegin()][0] + 1;
    if (k > 1 && !t01[u].empty()) try2 += dp[*t01[u].rbegin()][0] - dp[*t01[u].rbegin()][1] - 1;
    if (k > 1 && !s1[u].empty()) try2 += dp[*s1[u].rbegin()][1] + 2;
    dp[u][0] = max(try1, try2);
}
void Swappo(int u, int v) {
    Remove(u, v);
    Calc(u);
    Add(v, u);
    Calc(v);
}
void dfs(int u, int p) {
    for (int v: adj[u]) {
        if (v == p) continue;
        dfs(v, u);
        Add(u, v);
    }
    Calc(u);
}
void reroot(int u, int p) {
    ans = max({ans, dp[u][0], dp[u][1]});
    for (int v: adj[u]) {
        if (v == p) continue;
        Swappo(u, v);
        reroot(v, u);
        Swappo(v, u);
    }
}
signed main() {
    // #define filename "TKB"
    // freopen(filename".INP", "r", stdin);
    // freopen(filename".OUT", "w", stdout);
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int x, y;
        cin >> x >> y;
        adj[x].push_back(y);
        adj[y].push_back(x);
    }
    dfs(1, 1);
    for (int i = 1; i <= n; i++) cout << dp[i][0] << ' ' << dp[i][1] << '\n';
    // reroot(1, 1);
    cout << ans;
}
```