mech_combinatorics/
n_choose_k.rs

1use mech_core::*;
2use mech_core::value::ToUsize;
3#[cfg(feature = "matrix")]
4use mech_core::structures::matrix::Matrix;
5
6use std::fmt::Debug;
7use std::ops::{Add, AddAssign, Sub, Div};
8use num_traits::{Zero, One};
9use itertools::Itertools;
10use paste::paste;
11
12// Combinatorics N Choose K----------------------------------------------------
13
14#[derive(Debug)]
15pub struct NChooseK<T> {
16  n: Ref<T>,
17  k: Ref<T>,
18  out: Ref<T>,
19}
20impl<T> MechFunctionFactory for NChooseK<T> 
21where
22  T: Copy + Debug + Clone + Sync + Send + 'static +
23      Add<Output = T> + AddAssign +
24      Sub<Output = T> + Div<Output = T> +
25      Zero + One +
26      ConstElem + CompileConst + AsValueKind +
27      PartialEq + PartialOrd,
28  Ref<T>: ToValue,
29{
30  fn new(args: FunctionArgs) -> MResult<Box<dyn MechFunction>> {
31    match args {
32      FunctionArgs::Binary(out, arg1, arg2) => {
33        let n: Ref<T> = unsafe{ arg1.as_unchecked().clone() };
34        let k: Ref<T> = unsafe{ arg2.as_unchecked().clone() };
35        let out: Ref<T> = unsafe{ out.as_unchecked().clone() };
36        Ok(Box::new(Self{n, k, out}))
37      }
38      _ => Err(MechError{file: file!().to_string(), tokens: vec![], msg: "".to_string(), id: line!(), kind: MechErrorKind::IncorrectNumberOfArguments}),
39    }
40  }
41}
42impl<T> MechFunctionImpl for NChooseK<T>
43where
44  T: Copy + Debug + Clone + Sync + Send + 'static +
45      Add<Output = T> + AddAssign +
46      Sub<Output = T> + Div<Output = T> +
47      Zero + One +
48      PartialEq + PartialOrd,
49  Ref<T>: ToValue,
50{
51  fn solve(&self) {
52    let n_ptr = self.n.as_ptr();
53    let k_ptr = self.k.as_ptr();
54    let out_ptr = self.out.as_mut_ptr();
55    unsafe {
56      let n = *n_ptr;
57      let k = *k_ptr;
58      if k > n {
59        *out_ptr = T::zero(); // undefined for k > n
60        return;
61      }
62      let mut result = T::one();
63      let mut i = T::zero();
64      while i < k {
65        let numerator = n - i;
66        let denominator = i + T::one();
67        result = result * numerator / denominator;
68        i = i + T::one();
69      }
70      *out_ptr = result;
71    }
72  }
73  fn out(&self) -> Value {self.out.to_value()}
74  fn to_string(&self) -> String {format!("{:#?}", self)}
75}
76#[cfg(feature = "compiler")]
77impl<T> MechFunctionCompiler for NChooseK<T> 
78where
79    T: ConstElem + CompileConst + AsValueKind
80{
81  fn compile(&self, ctx: &mut CompileCtx) -> MResult<Register> {
82    let name = format!("NChooseK<{}>", T::as_value_kind());
83    compile_binop!(name, self.out, self.n, self.k, ctx, FeatureFlag::Custom(hash_str("combinatorics/n-choose-k")));
84  }
85}
86register_fxn_descriptor!(NChooseK, u8, "u8", i8, "i8", u16, "u16", i16, "i16", u32, "u32", i32, "i32", u64, "u64", i64, "i64", u128, "u128", i128, "i128", F32, "f32", F64, "f64", R64, "r64", C64, "c64");
87
88#[cfg(feature = "matrix")]
89#[derive(Debug)]
90pub struct NChooseKMatrix<T> {
91  n: Ref<Matrix<T>>,
92  k: Ref<T>,
93  out: Ref<Matrix<T>>,
94}
95#[cfg(feature = "matrix")]
96impl<T> MechFunctionFactory for NChooseKMatrix<T>
97where
98    T: Copy + Debug + Clone + Sync + Send + 'static +
99       ToUsize + std::fmt::Display +
100       Add<Output = T> + AddAssign +
101       Sub<Output = T> + Div<Output = T> +
102       Zero + One +
103      ConstElem + CompileConst + AsValueKind +
104       PartialEq + PartialOrd + ToMatrix,
105    Ref<T>: ToValue,
106    Matrix<T>: ToValue,
107{
108  fn new(args: FunctionArgs) -> MResult<Box<dyn MechFunction>> {
109    match args {
110      FunctionArgs::Binary(out, arg1, arg2) => {
111        let n: Ref<Matrix<T>> = unsafe{ arg1.as_unchecked().clone() };
112        let k: Ref<T> = unsafe{ arg2.as_unchecked().clone() };
113        let out: Ref<Matrix<T>> = unsafe{ out.as_unchecked().clone() };
114        Ok(Box::new(Self{n, k, out}))
115      }
116      _ => Err(MechError{file: file!().to_string(), tokens: vec![], msg: "".to_string(), id: line!(), kind: MechErrorKind::IncorrectNumberOfArguments}),
117    }
118  }
119}
120#[cfg(feature = "matrix")]
121impl<T> MechFunctionImpl for NChooseKMatrix<T>
122where
123    T: Copy + Debug + Clone + Sync + Send + 'static +
124       ToUsize + std::fmt::Display +
125       Add<Output = T> + AddAssign +
126       Sub<Output = T> + Div<Output = T> +
127       Zero + One +
128       PartialEq + PartialOrd + ToMatrix,
129    Ref<T>: ToValue,
130    Matrix<T>: ToValue,
131{
132  fn solve(&self) {
133      let n_matrix = self.n.borrow();
134      let k_scalar = *self.k.borrow();
135      let elements: Vec<T> = n_matrix.as_vec();
136      let k_usize: usize = k_scalar.to_usize();
137      // Check if k is greater than the number of elements. If it is, return an empty matrix.
138      if k_usize > elements.len() {
139          let empty = T::to_matrix(vec![], 0, k_usize);
140          *self.out.borrow_mut() = empty;
141          return;
142      }
143      // Generate combinations
144      let combinations: Vec<Vec<T>> = elements.iter().copied().combinations(k_usize).collect();
145      
146      // Reshape into output matrix
147      let rows = combinations.len();
148      let cols = k_usize;
149      let flat_data: Vec<T> = combinations.into_iter().flatten().collect();
150      let result = T::to_matrix(flat_data, cols, rows);
151      *self.out.borrow_mut() = result;
152  }
153  fn out(&self) -> Value { (*self.out.borrow()).to_value() }
154  fn to_string(&self) -> String { format!("{:#?}", self) }
155}
156#[cfg(all(feature = "matrix", feature = "compiler"))]
157impl<T> MechFunctionCompiler for NChooseKMatrix<T> 
158where
159    T: ConstElem + CompileConst + AsValueKind,
160{
161  fn compile(&self, ctx: &mut CompileCtx) -> MResult<Register> {
162    let name = format!("NChooseKMatrix<{}>", T::as_value_kind());
163    compile_binop!(name, self.out, self.n, self.k, ctx, FeatureFlag::Custom(hash_str("combinatorics/n-choose-k")));
164  }
165}
166register_fxn_descriptor!(NChooseKMatrix, u8, "u8", i8, "i8", u16, "u16", i16, "i16", u32, "u32", i32, "i32", u64, "u64", i64, "i64", u128, "u128", i128, "i128", F32, "f32", F64, "f64", R64, "r64", C64, "c64");
167
168
169fn impl_combinatorics_n_choose_k_fxn(n: Value, k: Value) -> Result<Box<dyn MechFunction>, MechError> {
170  match (n,k) {
171    #[cfg(feature = "u8")]
172    (Value::U8(n), Value::U8(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(u8::default())})),
173    #[cfg(feature = "u16")]
174    (Value::U16(n), Value::U16(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(u16::default())})),
175    #[cfg(feature = "u32")]
176    (Value::U32(n), Value::U32(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(u32::default())})),
177    #[cfg(feature = "u64")]
178    (Value::U64(n), Value::U64(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(u64::default())})),
179    #[cfg(feature = "u128")]
180    (Value::U128(n), Value::U128(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(u128::default())})),
181    #[cfg(feature = "i8")]
182    (Value::I8(n), Value::I8(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(i8::default())})),
183    #[cfg(feature = "i16")]
184    (Value::I16(n), Value::I16(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(i16::default())})),
185    #[cfg(feature = "i32")]
186    (Value::I32(n), Value::I32(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(i32::default())})),
187    #[cfg(feature = "i64")]
188    (Value::I64(n), Value::I64(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(i64::default())})),
189    #[cfg(feature = "i128")]
190    (Value::I128(n), Value::I128(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(i128::default())})),
191    #[cfg(feature = "f32")]
192    (Value::F32(n), Value::F32(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(F32::default())})),
193    #[cfg(feature = "f64")]
194    (Value::F64(n), Value::F64(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(F64::default())})),
195    #[cfg(feature = "rational")]
196    (Value::R64(n), Value::R64(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(R64::default())})),
197    #[cfg(feature = "complex")]
198    (Value::C64(n), Value::C64(k)) => Ok(Box::new(NChooseK{n: n, k: k, out: Ref::new(C64::default())})),
199    #[cfg(all(feature = "matrix", feature = "u8"))]
200    (Value::MatrixU8(n), Value::U8(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(u8::to_matrix(vec![], 0, 0))})),
201    #[cfg(all(feature = "matrix", feature = "u16"))]
202    (Value::MatrixU16(n), Value::U16(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(u16::to_matrix(vec![], 0, 0))})),
203    #[cfg(all(feature = "matrix", feature = "u32"))]
204    (Value::MatrixU32(n), Value::U32(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(u32::to_matrix(vec![], 0, 0))})),
205    #[cfg(all(feature = "matrix", feature = "u64"))]
206    (Value::MatrixU64(n), Value::U64(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(u64::to_matrix(vec![], 0, 0))})),
207    #[cfg(all(feature = "matrix", feature = "u128"))]
208    (Value::MatrixU128(n), Value::U128(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(u128::to_matrix(vec![], 0, 0))})),
209    #[cfg(all(feature = "matrix", feature = "i8"))]
210    (Value::MatrixI8(n), Value::I8(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(i8::to_matrix(vec![], 0, 0))})),
211    #[cfg(all(feature = "matrix", feature = "i16"))]
212    (Value::MatrixI16(n), Value::I16(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(i16::to_matrix(vec![], 0, 0))})),
213    #[cfg(all(feature = "matrix", feature = "i32"))]
214    (Value::MatrixI32(n), Value::I32(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(i32::to_matrix(vec![], 0, 0))})),
215    #[cfg(all(feature = "matrix", feature = "i64"))]
216    (Value::MatrixI64(n), Value::I64(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(i64::to_matrix(vec![], 0, 0))})),
217    #[cfg(all(feature = "matrix", feature = "i128"))]
218    (Value::MatrixI128(n), Value::I128(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(i128::to_matrix(vec![], 0, 0))})),
219    #[cfg(all(feature = "matrix", feature = "f32"))]
220    (Value::MatrixF32(n), Value::F32(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(F32::to_matrix(vec![], 0, 0))})),
221    #[cfg(all(feature = "matrix", feature = "f64"))]
222    (Value::MatrixF64(n), Value::F64(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(F64::to_matrix(vec![], 0, 0))})),
223    #[cfg(all(feature = "matrix", feature = "rational"))]
224    (Value::MatrixR64(n), Value::R64(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(R64::to_matrix(vec![], 0, 0))})),
225    #[cfg(all(feature = "matrix", feature = "complex"))]
226    (Value::MatrixC64(n), Value::C64(k)) => Ok(Box::new(NChooseKMatrix{n: Ref::new(n), k, out: Ref::new(C64::to_matrix(vec![], 0, 0))})),
227    x => Err(MechError{file: file!().to_string(), tokens: vec![], msg: format!("{:?}",x), id: line!(), kind: MechErrorKind::UnhandledFunctionArgumentKind }),
228  }
229}
230 
231pub struct CombinatoricsNChooseK {}
232impl NativeFunctionCompiler for CombinatoricsNChooseK {
233  fn compile(&self, arguments: &Vec<Value>) -> MResult<Box<dyn MechFunction>> {
234    if arguments.len() != 2 {
235      return Err(MechError{file: file!().to_string(), tokens: vec![], msg: "".to_string(), id: line!(), kind: MechErrorKind::IncorrectNumberOfArguments});
236    }
237    let n = arguments[0].clone();
238    let k = arguments[1].clone();
239
240    match impl_combinatorics_n_choose_k_fxn(n.clone(),k.clone()) {
241      Ok(fxn) => Ok(fxn),
242      Err(_) => {
243        match (n,k) {
244          (Value::MutableReference(n),Value::MutableReference(k)) => {let n_brrw = n.borrow();let k_brrw = k.borrow();impl_combinatorics_n_choose_k_fxn(n_brrw.clone(),k_brrw.clone())}
245          (n,Value::MutableReference(k)) => {let k_brrw = k.borrow(); impl_combinatorics_n_choose_k_fxn(n.clone(),k_brrw.clone())}
246          (Value::MutableReference(n),k) => {let n_brrw = n.borrow();impl_combinatorics_n_choose_k_fxn(n_brrw.clone(),k.clone())}
247          x => Err(MechError{file: file!().to_string(), tokens: vec![], msg: format!("{:?}",x), id: line!(), kind: MechErrorKind::UnhandledFunctionArgumentKind }),
248        }
249      }
250    }
251  }
252}
253
254register_descriptor! {
255  FunctionCompilerDescriptor {
256    name: "combinatorics/n-choose-k",
257    ptr: &CombinatoricsNChooseK{},
258  }
259}
mech_combinatorics/n_choose_k.rs

mech_combinatorics/
n_choose_k.rs
