Find Array Given Subset Sums

find-array-given-subset-sums.rs

/*
You are given an integer n representing the length of an unknown array that you are trying to recover. You are also given an array sums containing the values of all 2n subset sums of the unknown array (in no particular order).

Return the array ans of length n representing the unknown array. If multiple answers exist, return any of them.

An array sub is a subset of an array arr if sub can be obtained from arr by deleting some (possibly zero or all) elements of arr. The sum of the elements in sub is one possible subset sum of arr. The sum of an empty array is considered to be 0.

Note: Test cases are generated such that there will always be at least one correct answer.

Constraints:

1 <= n <= 15
sums.length == 2^n
-10^4 <= sums[i] <= 10^4
*/
use std::collections::HashMap;
impl Solution {
    pub fn recover_array(n: i32, sums: Vec<i32>) -> Vec<i32> {
        let mut count = HashMap::new();
        let mut min_sum = 10i32.pow(4);
        let mut max_sum = -(10i32.pow(4));

        for sum in sums {
            *count.entry(sum).or_insert(0) += 1;
            min_sum = min_sum.min(sum);
            max_sum = max_sum.max(sum);
        }

        let mut subset_sums = vec![];
        let mut sum_count = HashMap::new();
        let mut keys_sorted: Vec<&i32> = count.keys().collect();
        keys_sorted.sort_unstable();

        for key in keys_sorted {
            let value = key - min_sum;
            while sum_count.get(&value).unwrap_or(&0) < count.get(key).unwrap() {
                subset_sums.push(value);
                let mut new_sum_count = sum_count.clone();
                for (k, v) in sum_count.iter() {
                    *new_sum_count.entry(k + value).or_insert(0) += v;
                }
                sum_count = new_sum_count;
                if subset_sums.len() == n as usize {
                    break;
                }
            }
        }

        subset_sums.reverse();

        let mut pos_sum = 0;
        let mut neg_sum = 0;
        let mut ans = vec![];

        let found = find_valid(
            0,
            &mut pos_sum,
            &mut neg_sum,
            &subset_sums,
            max_sum,
            min_sum,
            &mut ans,
        );

        if found {
            return ans;
        }

        vec![]
    }
}

fn find_valid(
    index: usize,
    pos_sum: &mut i32,
    neg_sum: &mut i32,
    subset_sums: &[i32],
    max_sum: i32,
    min_sum: i32,
    ans: &mut Vec<i32>,
) -> bool {
    if index == subset_sums.len() {
        return true;
    }

    let value = subset_sums[index];

    for sign in &[1, -1] {
        let signed_value = value * sign;
        ans.push(signed_value);
        *pos_sum += signed_value.max(0);
        *neg_sum += signed_value.min(0);

        let found = find_valid(
            index + 1,
            pos_sum,
            neg_sum,
            subset_sums,
            max_sum,
            min_sum,
            ans,
        );

        if *pos_sum <= max_sum && *neg_sum >= min_sum && found {
            return true;
        }

        ans.pop();

        *pos_sum -= signed_value.max(0);
        *neg_sum -= signed_value.min(0);
    }

    false
}

struct Solution;