Maximum profit from trading stocks with discounts

impl Solution {
  pub fn max_profit(
    n: i32,
    present: Vec<i32>,
    future: Vec<i32>,
    hierarchy: Vec<Vec<i32>>,
    budget: i32,
  ) -> i32 {
    let n = n as usize;
    let budget = budget as usize;

    // Build tree (1-indexed -> 0-indexed)
    let mut children: Vec<Vec<usize>> = vec![vec![]; n];
    for e in &hierarchy {
      let u = (e[0] - 1) as usize;
      let v = (e[1] - 1) as usize;
      children[u].push(v);
    }

    // Tree DP. For each node, dp[node][cost] = max profit using exactly `cost` budget
    // considering the subtree of node, given whether the parent bought (discount flag).
    // dp_with[node][cost] and dp_without[node][cost]
    // Merge children one by one via knapsack.

    // Returns (dp_no_disc, dp_disc) where dp[cost] = max profit with `cost` spent
    fn dfs(
      v: usize,
      children: &Vec<Vec<usize>>,
      present: &Vec<i32>,
      future: &Vec<i32>,
      budget: usize,
    ) -> (Vec<i64>, Vec<i64>) {
      // dp_no_disc[cost]: max profit from subtree of v, spending exactly `cost`, when parent did NOT buy
      //                   => v can be bought at full price present[v]
      // dp_disc[cost]:   max profit from subtree of v, spending exactly `cost`, when parent DID buy
      //                   => v can be bought at price floor(present[v]/2)
      // Children get discount iff WE (v) buy.
      let price_full = present[v] as usize;
      let price_disc = (present[v] / 2) as usize;
      let profit_full = (future[v] - present[v]) as i64;
      let profit_disc = future[v] as i64 - price_disc as i64;

      let mut child_results: Vec<(Vec<i64>, Vec<i64>)> = vec![];
      for &c in &children[v] {
        child_results.push(dfs(c, children, present, future, budget));
      }

      // Merge children when v is NOT bought (children get no discount) -> comb_nodisc
      // Merge children when v IS bought (children get discount) -> comb_disc
      // Each child contributes either 0 (skip) or their subtree cost.
      // We merge incrementally.
      let mut comb_disc = vec![0i64; budget + 1];
      let mut comb_nodisc = vec![0i64; budget + 1];

      for (c_no, c_yes) in &child_results {
        // comb_disc (parent bought) merges with c_yes (child's parent bought)
        let mut new_disc = vec![i64::MIN; budget + 1];
        // comb_nodisc (parent not bought) merges with c_no (child's parent not bought)
        let mut new_nodisc = vec![i64::MIN; budget + 1];
        for ca in 0..=budget {
          if comb_disc[ca] < 0 { continue; }
          // don't buy anything from child subtree (cost 0, profit 0 already in children dp[0])
          // Actually c_yes[0] = 0 (buy nothing from child subtree - that's always valid)
          for cb in 0..=(budget - ca) {
            if c_yes[cb] >= 0 {
              let v2 = comb_disc[ca] + c_yes[cb];
              if new_disc[ca + cb] < v2 { new_disc[ca + cb] = v2; }
            }
          }
        }
        for ca in 0..=budget {
          if comb_nodisc[ca] < 0 { continue; }
          for cb in 0..=(budget - ca) {
            if c_no[cb] >= 0 {
              let v2 = comb_nodisc[ca] + c_no[cb];
              if new_nodisc[ca + cb] < v2 { new_nodisc[ca + cb] = v2; }
            }
          }
        }
        // Convert MIN back to 0 for valid no-spend option
        if new_disc[0] < 0 { new_disc[0] = 0; }
        if new_nodisc[0] < 0 { new_nodisc[0] = 0; }
        comb_disc = new_disc;
        comb_nodisc = new_nodisc;
      }

      // Build dp for v based on whether v is bought or not
      // dp_no_disc: parent didn't buy v => v costs price_full
      let mut dp_no_disc = vec![0i64; budget + 1];
      for c in 0..=budget {
        // don't buy v: use comb_nodisc
        if comb_nodisc[c] >= 0 {
          dp_no_disc[c] = dp_no_disc[c].max(comb_nodisc[c]);
        }
        // buy v at full price: use comb_disc for children
        if c >= price_full && comb_disc[c - price_full] >= 0 {
          dp_no_disc[c] = dp_no_disc[c].max(profit_full + comb_disc[c - price_full]);
        }
      }

      // dp_disc: parent bought v => v costs price_disc
      let mut dp_disc = vec![0i64; budget + 1];
      for c in 0..=budget {
        // don't buy v
        if comb_nodisc[c] >= 0 {
          dp_disc[c] = dp_disc[c].max(comb_nodisc[c]);
        }
        // buy v at discounted price
        if c >= price_disc && comb_disc[c - price_disc] >= 0 {
          dp_disc[c] = dp_disc[c].max(profit_disc + comb_disc[c - price_disc]);
        }
      }

      (dp_no_disc, dp_disc)
    }

    let (dp, _) = dfs(0, &children, &present, &future, budget);
    *dp.iter().max().unwrap() as i32
  }
}