diff --git a/src/lib.rs b/src/lib.rs index 14171f8..e9c92c6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -35,4 +35,5 @@ pub mod long_pressed_name; pub mod can_win_nim; pub mod hamming_distance; pub mod word_pattern; -pub mod update_matrix; \ No newline at end of file +pub mod update_matrix; +pub mod prison_cells_after_n_days; \ No newline at end of file diff --git a/src/prison_cells_after_n_days.rs b/src/prison_cells_after_n_days.rs new file mode 100644 index 0000000..7d7d0c3 --- /dev/null +++ b/src/prison_cells_after_n_days.rs @@ -0,0 +1,33 @@ +pub fn prison_after_n_days(cells: Vec, n: i32) -> Vec { + let mut cells = cells; + let mut start = (n - 1) % 14 + 1; + while start > 0 { + start -= 1; + let mut next = vec![0; 8]; + for i in 1..7 { + next[i] = if cells[i - 1] == cells[i + 1] { 1 } else { 0 }; + } + cells = next; + } + cells +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn example_1() { + assert_eq!( + prison_after_n_days(vec![0, 1, 0, 1, 1, 0, 0, 1], 7), + vec![0, 0, 1, 1, 0, 0, 0, 0] + ); + } + + #[test] + fn example_2() { + assert_eq!( + prison_after_n_days(vec![1, 0, 0, 1, 0, 0, 1, 0], 1000000000), + vec![0, 0, 1, 1, 1, 1, 1, 0] + ); + } +}