| 870 | |
| 871 | #[pyfunction] |
| 872 | fn comb(n: ArgIndex, k: ArgIndex, vm: &VirtualMachine) -> PyResult<BigInt> { |
| 873 | let n_int = n.into_int_ref(); |
| 874 | let n_big = n_int.as_bigint(); |
| 875 | let k_int = k.into_int_ref(); |
| 876 | let k_big = k_int.as_bigint(); |
| 877 | |
| 878 | if n_big.is_negative() { |
| 879 | return Err(vm.new_value_error("n must be a non-negative integer")); |
| 880 | } |
| 881 | if k_big.is_negative() { |
| 882 | return Err(vm.new_value_error("k must be a non-negative integer")); |
| 883 | } |
| 884 | |
| 885 | // Fast path: n fits in i64 |
| 886 | if let Some(ni) = n_big.to_i64() |
| 887 | && ni >= 0 |
| 888 | { |
| 889 | // k overflow or k > n means result is 0 |
| 890 | let ki = match k_big.to_i64() { |
| 891 | Some(k) if k >= 0 && k <= ni => k, |
| 892 | _ => return Ok(BigInt::from(0u8)), |
| 893 | }; |
| 894 | // Apply symmetry: use min(k, n-k) |
| 895 | let ki = ki.min(ni - ki); |
| 896 | if ki > 1 { |
| 897 | let result = pymath::math::integer::comb(ni, ki) |
| 898 | .map_err(|_| vm.new_value_error("comb() error"))?; |
| 899 | return Ok(result.into()); |
| 900 | } |
| 901 | // ki <= 1 cases |
| 902 | if ki == 0 { |
| 903 | return Ok(BigInt::from(1u8)); |
| 904 | } |
| 905 | return Ok(n_big.clone()); // ki == 1 |
| 906 | } |
| 907 | |
| 908 | // BigInt path: n doesn't fit in i64 |
| 909 | // Apply symmetry: k = min(k, n - k) |
| 910 | let n_minus_k = n_big - k_big; |
| 911 | if n_minus_k.is_negative() { |
| 912 | return Ok(BigInt::from(0u8)); |
| 913 | } |
| 914 | let effective_k = if &n_minus_k < k_big { |
| 915 | &n_minus_k |
| 916 | } else { |
| 917 | k_big |
| 918 | }; |
| 919 | |
| 920 | // k must fit in u64 |
| 921 | let ki: u64 = match effective_k.to_u64() { |
| 922 | Some(k) => k, |
| 923 | None => { |
| 924 | return Err( |
| 925 | vm.new_overflow_error(format!("min(n - k, k) must not exceed {}", u64::MAX)) |
| 926 | ); |
| 927 | } |
| 928 | }; |
| 929 | |