Add problem 3179: Find the N-th Value After K Seconds

Author: EFanZh

src/lib.rs

@@ -2123,6 +2123,7 @@ pub mod problem_3170_lexicographically_minimum_string_after_removing_stars;
 pub mod problem_3174_clear_digits;
 pub mod problem_3175_find_the_first_player_to_win_k_games_in_a_row;
 pub mod problem_3178_find_the_child_who_has_the_ball_after_k_seconds;
+pub mod problem_3179_find_the_n_th_value_after_k_seconds;
 pub mod problem_3350_adjacent_increasing_subarrays_detection_ii;
 
 #[cfg(test)]

src/problem_3179_find_the_n_th_value_after_k_seconds/mathematical.rs

@@ -0,0 +1,80 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl Solution {
+    const MODULUS: u64 = 1_000_000_007;
+
+    fn exp_mod(mut base: u64, mut exponent: u64) -> u64 {
+        let mut result = 1;
+
+        loop {
+            if exponent & 1 != 0 {
+                result = (result * base) % Self::MODULUS;
+
+                if exponent == 1 {
+                    break;
+                }
+            }
+
+            exponent >>= 1;
+            base = (base * base) % Self::MODULUS;
+        }
+
+        result
+    }
+
+    // See <https://en.wikipedia.org/wiki/Modular_multiplicative_inverse#Using_Euler's_theorem>.
+    fn mod_inv(x: u64) -> u64 {
+        Self::exp_mod(x, Self::MODULUS - 2)
+    }
+
+    fn binomial(n: u64, k: u64) -> u64 {
+        let n_minus_k = n - k;
+        let (k, n_minus_k) = if n_minus_k < k { (n_minus_k, k) } else { (k, n_minus_k) };
+        let mut factorial = 1;
+
+        for i in 2..=k {
+            factorial = (factorial * i) % Self::MODULUS;
+        }
+
+        let factorial_k = factorial;
+
+        for i in k + 1..=n_minus_k {
+            factorial = (factorial * i) % Self::MODULUS;
+        }
+
+        let factorial_n_minus_k = factorial;
+
+        for i in n_minus_k + 1..=n {
+            factorial = (factorial * i) % Self::MODULUS;
+        }
+
+        let factorial_n = factorial;
+
+        (factorial_n * Self::mod_inv((factorial_k * factorial_n_minus_k) % Self::MODULUS)) % Self::MODULUS
+    }
+
+    pub fn value_after_k_seconds(n: i32, k: i32) -> i32 {
+        let n = u64::from(n.cast_unsigned());
+        let k = u64::from(k.cast_unsigned());
+
+        Self::binomial(n + k - 1, n - 1) as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn value_after_k_seconds(n: i32, k: i32) -> i32 {
+        Self::value_after_k_seconds(n, k)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3179_find_the_n_th_value_after_k_seconds/mod.rs

@@ -0,0 +1,18 @@
+pub mod mathematical;
+
+pub trait Solution {
+    fn value_after_k_seconds(n: i32, k: i32) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [((4, 5), 56), ((5, 3), 35)];
+
+        for ((n, k), expected) in test_cases {
+            assert_eq!(S::value_after_k_seconds(n, k), expected);
+        }
+    }
+}