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

mech_combinatorics/
n_choose_k.rs
