I combined my two solution functions into a single one to pass through the data exactly once!
use std::fs::File;
use std::io::Read;
fn solve(tree: &Vec) -> (usize, usize) {
let mut s1 : usize = 0;
let mut iter = tree.iter();
let mut cstack = Vec::::new();
let mut mstack = Vec::::new();
let mut values = Vec::::new();
let nchilds = iter.next().unwrap();
cstack.push((*nchilds, *nchilds));
mstack.push(*iter.next().unwrap());
loop {
// No more children to process
if cstack.is_empty() { break; }
// Get next child
let (rem, orig) = cstack.pop().unwrap();
let mut sum : usize = 0;
match (rem, orig) {
// Node had no children, sum their metadata.
(0, 0) => {
let n = mstack.pop().unwrap();
for _ in 0..n { sum += iter.next().unwrap(); }
s1 += sum;
values.push(sum);
},
// Node had children, collapse their values.
(0, _) => {
let l = values.len() - orig - 1;
let n = mstack.pop().unwrap();
for _ in 0..n {
let m = iter.next().unwrap();
s1 += m;
if *m > orig || *m == 0 { continue; }
else { sum += values[l+m]; }
}
for _ in 0..orig { let _ = values.pop(); }
values.push(sum);
},
// More children to process.
(_, _) => {
let nchilds = iter.next().unwrap();
cstack.push((rem - 1, orig));
cstack.push((*nchilds, *nchilds));
mstack.push(*iter.next().unwrap());
},
}
}
let s2 = values.pop().unwrap();
(s1, s2)
}
fn main() {
let mut file = File::open("inputs/8.txt").unwrap();
let mut buffer = String::new();
let _ = file.read_to_string(&mut buffer);
let mut tree = Vec::::new();
for s in buffer.as_str().trim_end().split(' ') {
let i = s.parse::().unwrap();
tree.push(i);
}
let (sum, val) = solve(&tree);
println!("Sum of metadata: {}", sum);
println!("Root value: {}", val);
}
