use ndarray::{s, Array2};
use num_traits::PrimInt;
use rand::distr::{Distribution, StandardUniform};
use rand::{Rng, RngExt};
#[derive(Copy, Clone)]
pub enum MongePrim {
ConstantRows,
ConstantCols,
UpperRightOnes,
LowerLeftOnes,
}
impl MongePrim {
pub fn to_matrix<T: PrimInt, R: Rng>(self, m: usize, n: usize, rng: &mut R) -> Array2<T>
where
StandardUniform: Distribution<T>,
{
let mut matrix = Array2::from_elem((m, n), T::zero());
if m == 0 || n == 0 {
return matrix;
}
match self {
MongePrim::ConstantRows => {
for mut row in matrix.rows_mut() {
if rng.random::<bool>() {
row.fill(T::one())
}
}
}
MongePrim::ConstantCols => {
for mut col in matrix.columns_mut() {
if rng.random::<bool>() {
col.fill(T::one())
}
}
}
MongePrim::UpperRightOnes => {
let i = rng.random_range(0..m + 1) as isize;
let j = rng.random_range(0..n + 1) as isize;
matrix.slice_mut(s![..i, -j..]).fill(T::one());
}
MongePrim::LowerLeftOnes => {
let i = rng.random_range(0..m + 1) as isize;
let j = rng.random_range(0..n + 1) as isize;
matrix.slice_mut(s![-i.., ..j]).fill(T::one());
}
}
matrix
}
}
pub fn random_monge_matrix<R: Rng, T: PrimInt>(m: usize, n: usize, rng: &mut R) -> Array2<T>
where
StandardUniform: Distribution<T>,
{
let monge_primitives = [
MongePrim::ConstantRows,
MongePrim::ConstantCols,
MongePrim::LowerLeftOnes,
MongePrim::UpperRightOnes,
];
let mut matrix = Array2::from_elem((m, n), T::zero());
for _ in 0..(m + n) {
let monge = monge_primitives[rng.random_range(0..monge_primitives.len())];
matrix = matrix + monge.to_matrix(m, n, rng);
}
matrix
}