Maximum capacity within budget

impl Solution {
  pub fn max_capacity(costs: Vec<i32>, capacity: Vec<i32>, budget: i32) -> i32 {
    let n = costs.len();

    // For each cost value, track the best capacity
    // Then use sorting + two pointers to find the best pair

    // Create (cost, capacity) pairs and sort by cost
    let mut machines: Vec<(i32, i32)> = costs.iter().zip(capacity.iter()).map(|(&c, &cap)| (c, cap)).collect();
    machines.sort();

    // Best single machine with cost < budget
    let mut ans = 0i32;
    for &(c, cap) in &machines {
      if c < budget {
        ans = ans.max(cap);
      }
    }

    // For two machines: sort by cost, use two pointers
    // We need costs[i] + costs[j] < budget, maximize capacity[i] + capacity[j]
    //
    // For each right pointer, all machines with cost < budget - costs[right] are candidates for left.
    // We want the maximum capacity among those candidates (excluding index right itself).
    //
    // Better approach: sort by cost. Two pointer from left and right.
    // But we need max capacity, not just any pair.
    //
    // Alternative: for each cost threshold c, precompute the max capacity among machines with cost <= c.
    // Then for each machine i, the partner must have cost < budget - costs[i].
    // The best partner has max capacity among machines with cost < budget - costs[i], excluding i.

    // max_cap_at_cost[c] = max capacity among machines with cost exactly c
    // Then prefix max over costs.

    let max_cost = *costs.iter().max().unwrap() as usize;
    // For each cost value, store top-2 capacities (to handle excluding self)
    let mut top2: Vec<(i32, i32)> = vec![(-1, -1); max_cost + 1]; // (best, second_best)

    for i in 0..n {
      let c = costs[i] as usize;
      let cap = capacity[i];
      if cap >= top2[c].0 {
        top2[c].1 = top2[c].0;
        top2[c].0 = cap;
      } else if cap > top2[c].1 {
        top2[c].1 = cap;
      }
    }

    // prefix_top2[c] = top-2 capacities among machines with cost <= c
    let mut prefix_top2: Vec<(i32, i32)> = vec![(-1, -1); max_cost + 1];
    prefix_top2[0] = top2[0];
    for c in 1..=max_cost {
      prefix_top2[c] = prefix_top2[c - 1];
      // merge top2[c] into prefix_top2[c]
      for &cap in &[top2[c].0, top2[c].1] {
        if cap < 0 { continue; }
        if cap >= prefix_top2[c].0 {
          prefix_top2[c].1 = prefix_top2[c].0;
          prefix_top2[c].0 = cap;
        } else if cap > prefix_top2[c].1 {
          prefix_top2[c].1 = cap;
        }
      }
    }

    // For each machine i, find best partner
    for i in 0..n {
      let c = costs[i];
      let cap = capacity[i];
      let max_partner_cost = budget - c - 1; // strictly less than budget
      if max_partner_cost < 0 { continue; }
      let mpc = (max_partner_cost as usize).min(max_cost);

      let (best, second) = prefix_top2[mpc];
      // We need a partner different from machine i
      // But multiple machines can have same cost and capacity, so we use top-2
      // If best capacity is from a different machine, use it.
      // The issue: we can't tell if best is "us" or another machine with same cost/cap.
      //
      // Actually, we stored top-2 by capacity for each cost level, but the prefix merges them.
      // We don't track which machine they come from. Let's use a simpler approach:
      // the prefix_top2 tracks the two highest capacities among all machines with cost <= mpc.
      // If machine i's cost <= mpc and machine i has the top capacity, we should use second.

      // But this is tricky because multiple machines might share the same capacity.
      // Let's just check: if machine i's cost <= mpc (which it usually is unless mpc < c, but mpc = budget - c - 1)
      // Actually c + mpc = budget - 1, so if c <= mpc then 2c <= budget-1.
      // Machine i has cost c and cap. The partner must have cost <= mpc and be a different machine.

      // Since we track top-2 capacities (not indices), if best == cap and second < best,
      // and there's only one machine with that capacity at cost <= mpc, we should use second.
      // But we don't know if there's only one.

      // Safer approach: just check if we can use best. If best == cap, check if there's
      // another machine with the same capacity. We stored counts implicitly via top2 pairs.
      // Actually top2[c] = (best, second) where second < best only if there's one machine at best.
      // If two machines both have capacity = best, then top2[c] = (best, best).

      // For the prefix, same logic. If prefix_top2[mpc] = (X, X), there are at least 2 machines
      // with capacity X, so even if one is us, we can use another.

      if best < 0 { continue; }

      // Can we use `best`?
      // If cost[i] <= mpc, machine i is in the candidate pool
      if c as usize <= mpc {
        // machine i is in the pool, might be the best
        if cap == best {
          // Use second if available (second == best means there's another with same cap)
          if second >= 0 {
            ans = ans.max(cap + second);
          }
        } else {
          ans = ans.max(cap + best);
        }
      } else {
        // machine i is not in the pool (its cost > mpc), so best is definitely a different machine
        ans = ans.max(cap + best);
      }
    }

    ans
  }
}