Use this to learn the idea, then write your own version.
123 45class Solution {6public:7 int subtreeInversionSum(vector<vector<int>>& edges, vector<int>& nums, int k) {8 vector<vector<int>> adj(size(nums));9 const auto& iter_dfs = [&]() {10 using RET = vector<vector<int64_t>>;11 RET result{};12 vector<tuple<int, int, int, int, shared_ptr<RET>, RET *>> stk = {{1, 0, -1, -1, nullptr, &result}};13 while (!empty(stk)) {14 const auto [step, u, p, i, new_ret, ret] = stk.back(); stk.pop_back();15 if (step == 1) {16 ret->assign(2, vector<int64_t>(k, nums[u]));17 stk.emplace_back(4, -1, -1, -1, nullptr, ret);18 stk.emplace_back(2, u, p, 0, nullptr, ret);19 } else if (step == 2) {20 if (i == size(adj[u])) {21 continue;22 }23 const auto& v = adj[u][i];24 stk.emplace_back(2, u, p, i + 1, nullptr, ret);25 if (v == p) {26 continue;27 }28 const auto& new_ret = make_shared<RET>();29 stk.emplace_back(3, -1, -1, -1, new_ret, ret);30 stk.emplace_back(1, v, u, -1, nullptr, new_ret.get());31 } else if (step == 3) {32 auto& new_dp1 = (*new_ret)[0], &new_dp2 = (*new_ret)[1];33 auto& dp1 = (*ret)[0], &dp2 = (*ret)[1];34 for (int i = 0; i < k / 2; ++i) {35 dp1[i] = max(dp1[i] + new_dp1[(k - 2) - i], dp1[(k - 2) - i] + new_dp1[i]);36 dp2[i] = min(dp2[i] + new_dp2[(k - 2) - i], dp2[(k - 2) - i] + new_dp2[i]);37 }38 for (int i = k / 2; i < k; ++i) {39 dp1[i] += new_dp1[i];40 dp2[i] += new_dp2[i];41 }42 for (int i = k - 2; i >= 0; --i) {43 dp1[i] = max(dp1[i], dp1[i + 1]);44 dp2[i] = min(dp2[i], dp2[i + 1]);45 }46 } else if (step == 4) {47 auto& dp1 = (*ret)[0], &dp2 = (*ret)[1];48 const auto mx = max(dp1[0], -dp2[k - 1]);49 const auto mn = min(dp2[0], -dp1[k - 1]);50 dp1.insert(begin(dp1), mx); dp1.pop_back();51 dp2.insert(begin(dp2), mn); dp2.pop_back();52 }53 }54 return result[0][0];55 };56 57 for (auto& e : edges) {58 adj[e[0]].emplace_back(e[1]);59 adj[e[1]].emplace_back(e[0]);60 }61 return iter_dfs();62 }63};64 65666768class Solution2 {69public:70 int subtreeInversionSum(vector<vector<int>>& edges, vector<int>& nums, int k) {71 vector<vector<int>> adj(size(nums));72 const auto dfs = [&](this auto&& dfs, int u, int p) -> pair<vector<int64_t>, vector<int64_t>> {73 vector<int64_t> dp1(k, nums[u]), dp2(k, nums[u]);74 for (const auto& v : adj[u]) {75 if (v == p) {76 continue;77 }78 const auto& [new_dp1, new_dp2] = dfs(v, u);79 for (int i = 0; i < k / 2; ++i) {80 dp1[i] = max(dp1[i] + new_dp1[(k - 2) - i], dp1[(k - 2) - i] + new_dp1[i]);81 dp2[i] = min(dp2[i] + new_dp2[(k - 2) - i], dp2[(k - 2) - i] + new_dp2[i]);82 }83 for (int i = k / 2; i < k; ++i) {84 dp1[i] += new_dp1[i];85 dp2[i] += new_dp2[i];86 }87 for (int i = k - 2; i >= 0; --i) {88 dp1[i] = max(dp1[i], dp1[i + 1]);89 dp2[i] = min(dp2[i], dp2[i + 1]);90 }91 }92 const auto mx = max(dp1[0], -dp2[k - 1]);93 const auto mn = min(dp2[0], -dp1[k - 1]);94 dp1.insert(begin(dp1), mx); dp1.pop_back();95 dp2.insert(begin(dp2), mn); dp2.pop_back();96 return pair(dp1, dp2);97 };98 99 for (auto& e : edges) {100 adj[e[0]].emplace_back(e[1]);101 adj[e[1]].emplace_back(e[0]);102 }103 const auto& [dp1, _] = dfs(0, -1);104 return dp1[0];105 }106};107