Phrased in graph theoretic terms, this problem is asking us the following: We are given a rooted tree with each edge being labeled with a difficulty d and an enjoyment e. For any given query (s,c), we want to find the maximum total enjoyment path among all paths from the root satisfying the following condition: At most c of the edges can have difficulty greater than s.
Note that trying all vertices would take O(M∑δi), where δi is the depth of node i, which would take too long as it could be up to O(MN2).
Subtask 1 (N,M≤1000):
In this case, let Ei be the total enjoyment on the path starting from node i. We can compute Ei via the following dfs (all code snippets will assume we have shifted the nodes to be 0-indexed):
def precompute_E(p, e): E[0] = 0 for i in range(1, N): E[i] := E[p[i]] + e[i] return E
Similarly, we need to precompute the top 11 difficulty edges on the path from i. We can do this by first computing collecting all the difficulties, and checking the top 11.
def precompute_top_11(p, d): all_difficulties[0] = [] top_11_difficulties[0] = [] for i in range(1, N): all_weights[i] := all_weights[p[i]] + [d[i]] top_11_difficulties[i] = sorted(all_weights[i])[-11:] # In practice, we'll need to pad these
Then, for any query s, c, we need to find the maximum Ei over all nodes i whose c+1th difficulty is at most s.
Here is a C++ implementation. The precomputation takes O(N2logN) and processing queries takes O(N), giving a runtime of O(N2logN+MN).
#include <algorithm> #include <iostream> #include <vector> using namespace std; int main() { int N; cin >> N; int p[N], d[N], e[N]; for (int i = 1; i < N; i++) { cin >> p[i] >> d[i] >> e[i]; p[i]--; } long long E[N]; E[0] = 0; for (int i = 1; i < N; i++) { E[i] = E[p[i]] + e[i]; } vector<int> difficulties[N], top_11[N]; // pad with -1s, so that paths with k < 11 weights will // have their remaining 11-k elements filled with -1 difficulties[0] = vector<int>(11, -1); top_11[0] = vector<int>(11, -1); for (int i = 1; i < N; i++) { difficulties[i] = difficulties[p[i]]; difficulties[i].push_back(d[i]); top_11[i] = difficulties[i]; sort(top_11[i].begin(), top_11[i].end(), greater<int>()); top_11[i].resize(11); } int M; cin >> M; for (int i = 0; i < M; i++) { int S, C; cin >> S >> C; long long ans = 0; for (int j = 0; j < N; j++) { if (top_11[j][C] <= S) { ans = max(ans, E[j]); } } cout << ans << "\n"; } }
Alternatively, we can answer each query independently with a depth-first search in O(MN) time overall, though this doesn't naturally extend to the remaining subtasks.
Subtask 2 (c=0):
Here, we make two observations:
By preprocessing the array D and sorting by D, as well as maintaining cumulative maximum (i.e. precomputing max(E1,…,Ei) for each i), we can answer queries in O(logN) with O(NlogN) preprocessing, giving an O(NlogN+MlogN) runtime.
Here is a C++ implementation:
#include <algorithm> #include <iostream> #include <vector> using namespace std; int main() { int N; cin >> N; int p[N], d[N], e[N]; for (int i = 1; i < N; i++) { cin >> p[i] >> d[i] >> e[i]; p[i]--; } pair<int, long long> DE[N]; DE[0] = {-1, 0}; for (int i = 1; i < N; i++) { int Dpi = DE[p[i]].first; long long Epi = DE[p[i]].second; DE[i] = {max(Dpi, d[i]), Epi + e[i]}; } sort(DE, DE + N); for (int i = 1; i < N; i++) { DE[i].second = max(DE[i].second, DE[i-1].second); } int M, S, C; cin >> M; for (int i = 0; i < M; i++) { cin >> S >> C; int idx = lower_bound(DE, DE + N, make_pair(S+1, -1LL)) - DE; cout << DE[idx-1].second << "\n"; } }
Full Solution:
Note that observation 2 from Subtask 2 applies in general. If we can compute the top 11 difficulties array, then we can sort by cth difficulty for each c=1,2,…,11.
To compute the top 11 difficulties quickly, note that we can compute the top 11 difficulties for a given vertex i given the top 11 from pi and di. This is because anything difficulty that is not in the top 11 for pi won't be in the top 11 for i. This gives the following improvement:
def precompute_top_11(p, d): top_11_difficulties[0] = [] for i in range(1, N): relevant_weights = top_11_difficulties[p[i]] + [d[i]] top_11_difficulties[i] = sorted(relevant_weights)[-11:]
(Note: You can remove the sorting here by inserting di into the appropriate position. Depending on how you implement this, it will take either O(C) or O(logC) per step, which is a slight improvement from O(ClogC) per step shown here.)
Putting it all together, we get a solution that runs in O(NClogN+MlogN). Here is a C++ implementation:
#include <algorithm> #include <iostream> #include <vector> const int MAXC = 11; using namespace std; int main() { int N; cin >> N; int p[N], d[N], e[N]; for (int i = 1; i < N; i++) { cin >> p[i] >> d[i] >> e[i]; p[i]--; } long long E[N]; E[0] = 0; for (int i = 1; i < N; i++) { E[i] = E[p[i]] + e[i]; } vector<int> topC[N]; topC[0] = vector<int>(MAXC, -1); for (int i = 1; i < N; i++) { topC[i] = topC[p[i]]; topC[i].push_back(d[i]); sort(topC[i].begin(), topC[i].end(), greater<int>()); topC[i].resize(MAXC); } pair<int, long long> DE[MAXC][N]; for (int c = 0; c < MAXC; c++) { for (int i = 0; i < N; i++) { DE[c][i] = {topC[i][c], E[i]}; } sort(DE[c], DE[c] + N); for (int i = 1; i < N; i++) { DE[c][i].second = max(DE[c][i].second, DE[c][i-1].second); } } int M; cin >> M; for (int i = 0; i < M; i++) { int S, C; cin >> S >> C; int idx = lower_bound(DE[C], DE[C] + N, make_pair(S+1, -1LL)) - DE[C]; cout << DE[C][idx-1].second << endl; } }