#include "closing.h"
#include <iostream>
#include <vector>
#include <algorithm>
#define ll long long
using namespace std;
vector<vector<pair<int, int>>> pov;
vector<ll> distX, distY, cost1, cost2, level, type, q0;
vector<pair<ll, int>> q1, q2, q3;
void dfs(int u, int p, ll dist, vector<int>& distArray)
{
distArray[u] = dist;
for (pair<int, int> pr : pov[u]) {
int v = pr.first;
ll w = pr.second;
if (v == p) { continue; }
dfs(v, u, dist + w, distArray);
}
}
int max_score(int N, int X, int Y, ll K, vector <int> U, vector <int> V, vector <int> W)
{
pov.resize(N);
distX.resize(N,0);
distY.resize(N,0);
cost1.resize(N,0);
cost2.resize(N,0);
level.resize(N, 0);
type.resize(N, 0);
q0.resize(N);
q1.reserve(2 * N);
q2.reserve(N);
q3.reserve(N);
for (int j = 0; j < N - 1; j++)
{
pov[U[j]].push_back({ V[j], W[j] });
pov[V[j]].push_back({ U[j], W[j] });
}
dfs(X, X, 0, distX);
dfs(Y, Y, 0, distY);
for (int i = 0; i < N; i++)
{
cost1[i] = min(distX[i], distY[i]);
cost2[i] = max(distX[i], distY[i]) - cost1[i];
}
// case 1:
for (int i = 0; i < N; i++) q0[i] = -cost1[i];
sort(q0.begin(), q0.end());
ll budget = K;
int ans1 = 0;
while (q0.size() > 0)
{
if (q0.back() + budget >= 0)
{
budget += q0.back();
q0.pop_back();
ans1++;
}
else break;
}
// case 2:
budget = K;
int ans2 = 0;
for (int i = 0; i < N; i++) {
if (distX[i] + distY[i] == distY[X]) {
type[i] = 0;
budget -= cost1[i];
if (budget < 0) { return ans1; }
level[i] = 1;
q1.push_back({ -(cost2[i]), i });
ans2++;
}
else if (cost1[i] >= cost2[i]) { //MAYBE >=
type[i] = 1;
q2.push_back({ -(cost1[i] + cost2[i]), i });
q3.push_back({ -cost1[i], i });
}
else {
type[i] = 2;
q1.push_back({ -(cost1[i]), i });
q1.push_back({ -(cost2[i]), i });
}
}
if(q1.size())sort(q1.begin(), q1.end());
if(q2.size())sort(q2.begin(), q2.end());
if(q3.size())sort(q3.begin(), q3.end());
while (!q2.empty()) {
if (q2.back().first + budget < 0) { break; }
if (q1.size() >= 2) {
if (-q2.back().first <= -(q1.back().first + q1[q1.size() - 2].first)) {
budget += q2.back().first;
level[q2.back().second] = 2;
q2.pop_back();
ans2 += 2;
continue;
}
budget += q1.back().first;
level[q1.back().second]++;
q1.pop_back();
ans2 += 1;
continue;
}
budget += q2.back().first;
level[q2.back().second] = 2;
q2.pop_back();
ans2 += 2;
continue;
}
while (!q1.empty()) {
if (q1.back().first + budget < 0) { break; }
budget += q1.back().first;
level[q1.back().second]++;
q1.pop_back();
ans2 += 1;
}
while (!q3.empty()) {
if (q3.back().first + budget < 0) { break; }
if (level[q3.back().second] != 0) {
q3.pop_back();
continue;
}
budget += q3.back().first;
level[q3.back().second]++;
q3.pop_back();
ans2 += 1;
}
return max(ans1, ans2);
}
// int main() {
// cout << max_score(7, 0, 2, 10, { 0, 0, 1, 2, 2, 5 }, { 1, 3, 2, 4, 5, 6 }, { 2, 3, 4, 2, 5, 3 });
// return 0;
// }
Compilation message
closing.cpp: In function 'int max_score(int, int, int, long long int, std::vector<int>, std::vector<int>, std::vector<int>)':
closing.cpp:42:18: error: invalid initialization of reference of type 'std::vector<int>&' from expression of type 'std::vector<long long int>'
42 | dfs(X, X, 0, distX);
| ^~~~~
closing.cpp:12:46: note: in passing argument 4 of 'void dfs(int, int, long long int, std::vector<int>&)'
12 | void dfs(int u, int p, ll dist, vector<int>& distArray)
| ~~~~~~~~~~~~~^~~~~~~~~
closing.cpp:43:18: error: invalid initialization of reference of type 'std::vector<int>&' from expression of type 'std::vector<long long int>'
43 | dfs(Y, Y, 0, distY);
| ^~~~~
closing.cpp:12:46: note: in passing argument 4 of 'void dfs(int, int, long long int, std::vector<int>&)'
12 | void dfs(int u, int p, ll dist, vector<int>& distArray)
| ~~~~~~~~~~~~~^~~~~~~~~
