Add problem 2049: Count Nodes With the Highest Score

src/lib.rs

@@ -1494,6 +1494,7 @@ pub mod problem_2042_check_if_numbers_are_ascending_in_a_sentence;
 pub mod problem_2043_simple_bank_system;
 pub mod problem_2044_count_number_of_maximum_bitwise_or_subsets;
 pub mod problem_2047_number_of_valid_words_in_a_sentence;
+pub mod problem_2049_count_nodes_with_the_highest_score;
 pub mod problem_2053_kth_distinct_string_in_an_array;
 pub mod problem_2055_plates_between_candles;
 pub mod problem_2057_smallest_index_with_equal_value;

src/problem_2049_count_nodes_with_the_highest_score/mod.rs

@@ -0,0 +1,18 @@
+pub mod recursive;
+
+pub trait Solution {
+    fn count_highest_score_nodes(parents: Vec<i32>) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [(&[-1, 2, 0, 2, 0] as &[_], 3), (&[-1, 2, 0], 2), (&[-1, 0], 2)];
+
+        for (parents, expected) in test_cases {
+            assert_eq!(S::count_highest_score_nodes(parents.to_vec()), expected);
+        }
+    }
+}

src/problem_2049_count_nodes_with_the_highest_score/recursive.rs

@@ -0,0 +1,104 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::cmp::Ordering;
+
+struct Context<'a> {
+    graph: &'a [(u32, u32)],
+    max_score: u64,
+    max_score_count: u32,
+}
+
+impl Context<'_> {
+    fn update_score(&mut self, score: u64) {
+        match score.cmp(&self.max_score) {
+            Ordering::Less => {}
+            Ordering::Equal => self.max_score_count += 1,
+            Ordering::Greater => {
+                self.max_score = score;
+                self.max_score_count = 1;
+            }
+        }
+    }
+}
+
+impl Solution {
+    fn dfs(context: &mut Context, node: u32) -> u32 {
+        let n = context.graph.len() as u32;
+        let (left, right) = context.graph[node as usize];
+        let mut count = 1;
+
+        let score = if left == 0 {
+            u64::from(n - 1)
+        } else {
+            let left_nodes = Self::dfs(context, left);
+
+            count += left_nodes;
+
+            if right == 0 {
+                u64::from(n - 1 - left_nodes) * u64::from(left_nodes)
+            } else {
+                let right_nodes = Self::dfs(context, right);
+
+                count += right_nodes;
+
+                u64::from(n - 1 - left_nodes - right_nodes) * u64::from(left_nodes) * u64::from(right_nodes)
+            }
+        };
+
+        context.update_score(score);
+
+        count
+    }
+
+    pub fn count_highest_score_nodes(parents: Vec<i32>) -> i32 {
+        let mut graph = vec![(0_u32, 0_u32); parents.len()].into_boxed_slice();
+
+        for (node, &parent) in (1..).zip(&parents[1..]) {
+            let parent = &mut graph[parent as u32 as usize];
+
+            if parent.0 == 0 {
+                parent.0 = node;
+            } else {
+                parent.1 = node;
+            }
+        }
+
+        drop(parents);
+
+        let mut context = Context {
+            graph: &graph,
+            max_score: 0,
+            max_score_count: 0,
+        };
+
+        let (left, right) = graph[0];
+        let left_count = Self::dfs(&mut context, left);
+        let mut score = u64::from(left_count);
+
+        if right != 0 {
+            score *= u64::from(Self::dfs(&mut context, right));
+        };
+
+        context.update_score(score);
+
+        context.max_score_count as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn count_highest_score_nodes(parents: Vec<i32>) -> i32 {
+        Self::count_highest_score_nodes(parents)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}