CF1746D（记忆化搜索，DP，贪心）

CF1746D（记忆化搜索，DP，贪心）

https://codeforces.com/contest/1746/problem/d

题意

给一棵树，树上每个点有一个权值 \(s_i\), 有一个整数 \(k\)。表示从根节点出发的简单路径的数量。给出约束：对一点 \(u\) ，它的儿子所经过的简单路径的数量差不能超过1。

设每个点经过的简单路径的数量为 \(c_i\)。最大化 \(\sum c_i s_i\)。

思路

两个显然的结论，一定要让路径走到底，到达叶子才有可能最优。从某点经过向下的简单路径对它的儿子的分配数量一定是 \(\lfloor \frac{c_i}{son} \rfloor\) 或者是 \(\lfloor \frac{c_i}{son} \rfloor+ 1\)。

然后考虑到可以用 \(f(u, sum)\) 表示所有子问题的解，考虑记忆化搜索。

定义 \(f[u][sum]\) 表示当前点 \(u\) 分配路径数量为 \(sum\) 的答案。

转移就是所有儿子的 \(f[v][\frac{sum}{son}]\) 或者 \(f[v][\frac{sum}{son}+1]\) 的和。问题是如何选择这两种情况。最显然的一个解法是做一个01背包，显然时间上过不去。这里和普通背包不同的是限制了取 "1" 的数量，这时可以利用这一个条件进行贪心排序：先全部选 "0"，再贪心选择 "1" 进行替换。什么样的 "1" 最优呢：按照“0”和“1”的差值排序取最大即可。

于是可以完成\(O(sonlog(son))\) 的转移。

最后如果树是一条链，每次需要向两个情况搜索，且sum无法很快的变成1，算法复杂度将变成指数级的。所以需要特判掉链的情况。对于其他情况，由于 \(sum\) 很快会减小，树的深度最多向下 \(logn\) 层，所以复杂度是正确的。

#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
#include<set>
#include<bitset>
#include<queue>
#include<map>
#include<stack>
#include<string>
#include<random>
#include<cassert>
#include<functional>
#include<iomanip>
#define inf 0x3f3f3f3f
#define linf 0x3f3f3f3f3f3f3f3fll
#define endl '\n'
#define ll long long
#define int long long
#define SZ(x) (int)x.size()
#define rep(i,a,n) for(int i = (a);i <= (n);i++)
#define dec(i,n,a) for(int i = (n);i >= (a);i--)
using namespace std;
using PII = pair<int, int>;
mt19937 mrand(random_device{}());
int rnd(int x) { return mrand() % x;}
const int N =10 + 2e5 ,mod=1e9 + 7;

void solve()
{    
    int n, k; cin >> n >> k;
    vector<vector<int>> adj(n + 1);
    rep(i,1, n - 1) {
        int p; cin >> p;
        adj[p].emplace_back(i + 1);
    }
    vector<int> s(n + 1); rep(i, 1, n) cin >> s[i];

    map<int, map<int,int>> f;// f[][] = ?
    function<int(int, int)> dfs = [&](int u, int sum) {
        if(f[u].count(sum)) return f[u][sum];
        int ans = sum * s[u];
        
        if(!SZ(adj[u])) {
            return f[u][sum] = ans;
        }

        vector<PII> t;
        int o = sum % SZ(adj[u]);
        if(o) {
            for(auto v : adj[u]) {
                PII cur;
                cur.first = dfs(v, sum / SZ(adj[u]));
                cur.second = dfs(v, sum / SZ(adj[u]) + 1);
                t.push_back(cur);

                ans += cur.first;
            }

            sort(t.begin(), t.end(), [&](PII &A, PII &B) {
                return A.second - A.first > B.second - B.first;
            });

            rep(i, 0, o - 1) ans -= t[i].first, ans += t[i].second;
        }else {
            for(auto v : adj[u]) {
                ans += dfs(v, sum / SZ(adj[u]));
            }
        }
        return f[u][sum] = ans;
    };

    cout << dfs(1, k) << endl;
}
signed main()
{
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

    int T;cin>>T;
    while(T--)
        solve();

    return 0;
}