#include "closing.h"
#include <algorithm>
#include <iostream>
#include <numeric>
#include <cassert>
#include <vector>
#include <queue>
typedef long long llong;
const int MAXN = 500 + 10;
const llong INF = 8e18;
const int INTINF = 2e9;
int n;
int x;
int y;
llong k;
int sz[MAXN];
int par[MAXN];
bool isOnPath[MAXN];
int suffixSizes[MAXN];
std::vector <int> path;
std::vector <int> pathSizes;
std::vector <std::pair <int,int>> g[MAXN];
llong dpR[MAXN][2 * MAXN][2][2]; // root
llong dpC[MAXN][2 * MAXN][2][2]; // children, knapsack
llong dpS[MAXN][2 * MAXN][2][2]; // stick
bool blR[MAXN][2 * MAXN][2][2];
bool blC[MAXN][2 * MAXN][2][2];
bool blS[MAXN][2 * MAXN][2][2];
llong distX[MAXN];
llong distY[MAXN];
void reset()
{
path.clear();
pathSizes.clear();
for (int i = 0 ; i < n ; ++i)
{
sz[i] = 0;s
g[i].clear()
isOnPath[i] = false;
distX[i] = distY[i] = 0;
for (int j = 0 ; j <= 2 * n ; ++j)
{
for (int flagX = 0 ; flagX < 2 ; ++flagX)
{
for (int flagY = 0 ; flagY < 2 ; ++flagY)
{
dpR[i][j][flagX][flagY] = 0;
dpC[i][j][flagX][flagY] = 0;
dpS[i][j][flagX][flagY] = 0;
blR[i][j][flagX][flagY] = false;
blC[i][j][flagX][flagY] = false;
blS[i][j][flagX][flagY] = false;
}
}
}
}
}
void dfs(int node, int par, llong to[], llong dist)
{
to[node] = dist;
for (const auto &[u, d] : g[node])
{
if (u == par)
{
continue;
}
dfs(u, node, to, dist + d);
}
}
bool dfsPath(int node, int par)
{
if (node == y)
{
path.push_back(node);
return true;
}
for (const auto &[u, d] : g[node])
{
if (u == par)
{
continue;
}
if (dfsPath(u, node))
{
path.push_back(node);
return true;
}
}
return false;
}
void dfsSize(int node, int p)
{
sz[node] = 1;
par[node] = p;
for (const auto &[u, d] : g[node])
{
if (p == u)
{
continue;
}
dfsSize(u, node);
if (!isOnPath[u]) sz[node] += sz[u];
}
}
llong fRoot(int, int, int, int);
llong fChildren(int, int, int, int, int);
llong fStick(int, int, int, int);
llong fRoot(int node, int choose, int flagX, int flagY)
{
if (choose == 0)
{
return 0;
}
if (!flagX && !flagY)
{
return k + 1;
}
if (blR[node][choose][flagX][flagY])
{
return dpR[node][choose][flagX][flagY];
}
blR[node][choose][flagX][flagY] = true;
dpR[node][choose][flagX][flagY] = fChildren(node, 0, choose, flagX, flagY);
dpR[node][choose][flagX][flagY] = std::min(dpR[node][choose][flagX][flagY], k + 1);
return dpR[node][choose][flagX][flagY];
}
llong fChildren(int node, int idx, int choose, int flagX, int flagY)
{
if (idx == g[node].size())
{
if (choose == 0) return 0;
return k + 1;
}
int current = g[node][idx].first;
if (isOnPath[current] || current == par[node])
{
return fChildren(node, idx + 1, choose, flagX, flagY);
}
if (blC[current][choose][flagX][flagY])
{
return dpC[current][choose][flagX][flagY];
}
blC[current][choose][flagX][flagY] = true;
dpC[current][choose][flagX][flagY] = k + 1;
for (int currX = 0 ; currX <= flagX ; ++currX)
{
for (int currY = 0 ; currY <= flagY ; ++currY)
{
for (int subtreeChoose = std::max(choose - 2 * suffixSizes[(idx + 1 == g[node].size() ? 0 : g[node][idx + 1].first)], currX + currY) ; subtreeChoose <= std::min(2 * sz[g[node][idx].first], choose) ; ++subtreeChoose)
{
if (fChildren(node, idx + 1, choose - subtreeChoose, flagX, flagY) > k) continue;
dpC[current][choose][flagX][flagY] = std::min(dpC[current][choose][flagX][flagY], fChildren(node, idx + 1, choose - subtreeChoose, flagX, flagY) + std::max((currX == 0 ? 0 : distX[current]), (currY == 0 ? 0 : distY[current])) + fRoot(current, subtreeChoose - currX - currY, currX, currY));
}
}
}
return dpC[current][choose][flagX][flagY];
}
llong fStick(int idx, int choose, int flagX, int flagY)
{
if (idx == path.size())
{
if (choose == 0) return 0;
return k + 1;
}
if (blS[idx][choose][flagX][flagY])
{
return dpS[idx][choose][flagX][flagY];
}
blS[idx][choose][flagX][flagY] = true;
dpS[idx][choose][flagX][flagY] = k + 1;
for (int currX = 0 ; currX <= flagX ; ++currX)
{
for (int currY = flagY ; currY < 2 ; ++currY)
{
for (int subtreeChoose = std::max(choose - 2 * pathSizes[idx + 1], currX + currY) ; subtreeChoose <= std::min(2 * sz[path[idx]], choose) ; ++subtreeChoose)
{
if (fStick(idx + 1, choose - subtreeChoose, currX, currY) > k) continue;
dpS[idx][choose][flagX][flagY] = std::min(dpS[idx][choose][flagX][flagY], fStick(idx + 1, choose - subtreeChoose, currX, currY) + std::max((currX == 0 ? 0 : distX[path[idx]]), (currY == 0 ? 0 : distY[path[idx]])) + fRoot(path[idx], subtreeChoose - currX - currY, currX, currY));
}
}
}
return dpS[idx][choose][flagX][flagY];
}
int max_score(int N, int X, int Y, long long K, std::vector <int> U, std::vector <int> V, std::vector <int> W)
{
reset();
n = N;
x = X;
y = Y;
k = K;
for (int i = 0 ; i < n - 1 ; ++i)
{
g[U[i]].emplace_back(V[i], W[i]);
g[V[i]].emplace_back(U[i], W[i]);
}
dfs(x, -1, distX, 0);
dfs(y, -1, distY, 0);
dfsPath(x, -1);
dfsSize(x, -1);
for (int i = 0 ; i < n ; ++i)
{
for (int j = g[i].size() - 1 ; j >= 0 ; --j)
{
int add = 0;
suffixSizes[g[i][j].first] = suffixSizes[(j + 1 == g[i].size() ? 0 : g[i][j + 1].first)];
if (!isOnPath[g[i][j].first] && g[i][j].first != par[i])
{
add = sz[g[i][j].first];
}
suffixSizes[g[i][j].first] += add;
}
}
std::reverse(path.begin(), path.end());
for (const int &i : path)
{
isOnPath[i] = true;
}
pathSizes.resize(path.size() + 1, 0);
for (int i = path.size() - 1 ; i >= 0 ; --i)
{
pathSizes[i] = pathSizes[i + 1] + sz[path[i]];
}
for (int i = 2 * n ; i >= 1 ; --i)
{
if (fStick(0, i, 1, 0) <= k)
{
return i;
}
}
return 0;
}
Compilation message
closing.cpp: In function 'void reset()':
closing.cpp:40:19: error: 's' was not declared in this scope
40 | sz[i] = 0;s
| ^
closing.cpp: In function 'llong fChildren(int, int, int, int, int)':
closing.cpp:146:13: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
146 | if (idx == g[node].size())
| ~~~~^~~~~~~~~~~~~~~~~
closing.cpp:170:81: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
170 | for (int subtreeChoose = std::max(choose - 2 * suffixSizes[(idx + 1 == g[node].size() ? 0 : g[node][idx + 1].first)], currX + currY) ; subtreeChoose <= std::min(2 * sz[g[node][idx].first], choose) ; ++subtreeChoose)
| ~~~~~~~~^~~~~~~~~~~~~~~~~
closing.cpp: In function 'llong fStick(int, int, int, int)':
closing.cpp:183:13: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
183 | if (idx == path.size())
| ~~~~^~~~~~~~~~~~~~
closing.cpp: In function 'int max_score(int, int, int, long long int, std::vector<int>, std::vector<int>, std::vector<int>)':
closing.cpp:235:61: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
235 | suffixSizes[g[i][j].first] = suffixSizes[(j + 1 == g[i].size() ? 0 : g[i][j + 1].first)];
| ~~~~~~^~~~~~~~~~~~~~
