Add problem 1617: Count Subtrees With Max Distance Between Cities

src/lib.rs

@@ -1330,6 +1330,7 @@ pub mod problem_1611_minimum_one_bit_operations_to_make_integers_zero;
 pub mod problem_1614_maximum_nesting_depth_of_the_parentheses;
 pub mod problem_1615_maximal_network_rank;
 pub mod problem_1616_split_two_strings_to_make_palindrome;
+pub mod problem_1617_count_subtrees_with_max_distance_between_cities;
 pub mod problem_1619_mean_of_array_after_removing_some_elements;
 pub mod problem_1624_largest_substring_between_two_equal_characters;
 pub mod problem_1629_slowest_key;

src/problem_1617_count_subtrees_with_max_distance_between_cities/brute_force.rs

@@ -0,0 +1,120 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::convert::TryInto;
+
+enum ChildrenDepths {
+    None,
+    Single(u8),
+    Multiple(u8, u8),
+}
+
+impl Solution {
+    fn dfs(graph: &[u16; 16], visited: &mut u16, node: u8) -> (u8, u8) {
+        let mut children = graph[usize::from(node)];
+        let mut max_distance = 0;
+        let mut max_depth = 0;
+        let mut children_depths = ChildrenDepths::None;
+
+        while children != 0 {
+            let child = children.trailing_zeros() as _;
+
+            children &= children - 1;
+
+            let probe = 1_u16 << child;
+
+            if *visited & probe == 0 {
+                *visited |= probe;
+
+                let (child_max_distance, child_depth) = Self::dfs(graph, visited, child);
+
+                max_distance = max_distance.max(child_max_distance);
+                max_depth = max_depth.max(child_depth + 1);
+
+                children_depths = match children_depths {
+                    ChildrenDepths::None => ChildrenDepths::Single(child_depth),
+                    ChildrenDepths::Single(depth) => {
+                        if child_depth > depth {
+                            ChildrenDepths::Multiple(child_depth, depth)
+                        } else {
+                            ChildrenDepths::Multiple(depth, child_depth)
+                        }
+                    }
+                    ChildrenDepths::Multiple(existing_depth_1, existing_depth_2) => {
+                        if child_depth > existing_depth_2 {
+                            if child_depth > existing_depth_1 {
+                                ChildrenDepths::Multiple(child_depth, existing_depth_1)
+                            } else {
+                                ChildrenDepths::Multiple(existing_depth_1, child_depth)
+                            }
+                        } else {
+                            continue;
+                        }
+                    }
+                }
+            }
+        }
+
+        (
+            max_distance.max(match children_depths {
+                ChildrenDepths::None => 0,
+                ChildrenDepths::Single(depth) => depth + 1,
+                ChildrenDepths::Multiple(depth_1, depth_2) => depth_1 + depth_2 + 2,
+            }),
+            max_depth,
+        )
+    }
+
+    pub fn count_subgraphs_for_each_diameter(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
+        let n = n as u32 as usize;
+        let mut graph = [0_u16; 16];
+
+        for edge in edges {
+            let [from, to]: [_; 2] = edge.try_into().ok().unwrap();
+            let from = from - 1;
+            let to = to - 1;
+
+            graph[from as u32 as usize] |= 1 << to;
+            graph[to as u32 as usize] |= 1 << from;
+        }
+
+        let mut result = vec![0; n - 1];
+
+        for subset in 3_u16..(1 << n) {
+            if !subset.is_power_of_two() {
+                let mut subgraph = graph;
+
+                for neighbor in &mut subgraph[..n] {
+                    *neighbor &= subset;
+                }
+
+                let node = subset.trailing_zeros() as _;
+                let mut visited = 1_u16 << node;
+                let candidate_max_distance = Self::dfs(&subgraph, &mut visited, node).0;
+
+                if visited == subset {
+                    result[usize::from(candidate_max_distance) - 1] += 1;
+                }
+            }
+        }
+
+        result
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn count_subgraphs_for_each_diameter(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
+        Self::count_subgraphs_for_each_diameter(n, edges)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_1617_count_subtrees_with_max_distance_between_cities/mod.rs

@@ -0,0 +1,33 @@
+pub mod brute_force;
+
+pub trait Solution {
+    fn count_subgraphs_for_each_diameter(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((4, &[[1, 2], [2, 3], [2, 4]] as &[_]), &[3, 4, 0] as &[_]),
+            ((2, &[[1, 2]]), &[1]),
+            ((3, &[[1, 2], [2, 3]]), &[2, 1]),
+            (
+                (8, &[[1, 5], [2, 3], [2, 5], [2, 8], [4, 7], [6, 7], [6, 8]]),
+                &[7, 8, 8, 6, 5, 2, 0],
+            ),
+            (
+                (9, &[[1, 3], [1, 5], [1, 6], [2, 4], [3, 4], [5, 7], [5, 8], [6, 9]]),
+                &[8, 11, 16, 20, 11, 0, 0, 0],
+            ),
+        ];
+
+        for ((n, edges), expected) in test_cases {
+            assert_eq!(
+                S::count_subgraphs_for_each_diameter(n, edges.iter().copied().map(Vec::from).collect(),),
+                expected,
+            );
+        }
+    }
+}