Submission #67973027

---

Source Code Expand

#include <bits/stdc++.h>
#include <bits/extc++.h>
using namespace std;
using namespace __gnu_pbds;
using namespace __gnu_cxx;
const int maxn = 2e5+5;
long long a[maxn];
vector<int>G[maxn];
int n,k;
long long dp[maxn][7][3],dp2[maxn][7][3];

void dfs(int d,int fa)
{
    for (int i=1;i<=k;i++)
        dp[d][i][1]=a[d];
    for (auto i:G[d])
    {
        if (i!=fa)
        {
            dfs(i,d);
            for (int j=k;j>0;j--)
            {
                for (int kk=1;kk<=j;kk++)
                {
                    dp2[d][j][0]=max(dp2[d][j][0],dp[d][j-kk][0]+max({dp[i][kk][0],dp[i][kk][1],dp[i][kk][2]}));
                    dp2[d][j][2]=max(dp2[d][j][2],dp[d][j-kk][2]+max({dp[i][kk][0],dp[i][kk][1],dp[i][kk][2]}));
                    dp2[d][j][2]=max(dp2[d][j][2],dp[d][j-kk+1][1]+dp[i][kk][1]);
                    dp2[d][j][2]=max(dp2[d][j][2],dp[d][j-kk][1]+dp[i][kk+1][1]);
                    dp2[d][j][1]=max({dp2[d][j][1],dp[d][j-kk][1]+max({dp[i][kk][0],dp[i][kk][1],dp[i][kk][2]}),dp[d][j-kk][0]+dp[i][kk][1]+a[d]});
                }
            }
            for (int j=k;j>0;j--)
                for (int kk=0;kk<=2;kk++)
                    dp[d][j][kk]=max(dp[d][j][kk],dp2[d][j][kk]);
        }
    }
    for (int i=1;i<=k;i++)
        dp[d][i][1]=max(dp[d][i][1],dp[d][i-1][0]+a[d]);
}
signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    cin>>n>>k;
    for (int i=1;i<=n;i++)
        cin>>a[i];
    for (int i=1;i<n;i++)
    {
        int u,v;
        cin>>u>>v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1,0);
    long long ans=0;
    for (int i=0;i<=k;i++)
        for (int j=0;j<3;j++)
            ans=max(ans,dp[1][i][j]);
    cout<<ans;
}

Submission Info

Judge Result