diff --git a/modules/maths/math_functions.nu b/modules/maths/math_functions.nu index b85c478..7ec03e3 100644 --- a/modules/maths/math_functions.nu +++ b/modules/maths/math_functions.nu @@ -157,3 +157,84 @@ export def scale-minmax-table [a, b,input?] { ($x | column2 $i) | scale-minmax $a $b | wrap ($name_cols | get $i) } | reduce {|it, acc| $acc | merge {$it}} } + + +# compute the cartesian product of any number of lists +# +# basically, if you give `iter cartesian product` *n* lists, from *i_1* to *i_n*, +# it will compute recursively the cartesian product of the first one with the +# `iter cartesian product` of the rest, i.e. if we call CP the two-set cartesian +# product and ICP the multi cartesian product here, we have +# +# *ICP(i_1, i_2, ..., i_n) == CP(i_1, ICP(i_2, ..., i_n))* +# +# # Example +#```nushell +# use std ["assert equal" "iter iter cartesian product"] +# +# let res = ( +# iter cartesian product [1, 2] [3, 4] +# ) +# +# assert equal $res [ +# [1, 3], +# [1, 4], +# [2, 3], +# [2, 4], +# ] +# ``` +export def "cartesian product" [ + ...iters: list # the iterables you want the cartesian product of +]: nothing -> list> { + def aux [a: list>]: nothing -> list> { + if ($a | is-empty) { + return [] + } + + let head = $a | first + let tail = aux ($a | skip 1) + + if ($head | is-empty) { + return $tail + } else if ($tail | is-empty) { + return $head + } + + $head | each {|h| $tail | each {|t| [$h, $t]}} | flatten | each { flatten } + } + + aux $iters +} + + +#[test] +def cartesian_product [] { + use std assert + + # emptyness + assert equal (cartesian product [] []) [] + assert equal (cartesian product []) [] + assert equal (cartesian product) [] + + # symmetry + assert equal (cartesian product [1, 2] []) [1, 2] + assert equal (cartesian product [] [1, 2]) [1, 2] + # NOTE: `cartesian product` might not preserve the order of the elements produced + assert equal ( + cartesian product [1, 2] [3, 4] | each { sort } | sort + ) ( + cartesian product [3, 4] [1, 2] | each { sort } | sort + ) + + assert equal ( + cartesian product [1, 2] [3, 4] + ) [ + [1, 3], [1, 4], [2, 3], [2, 4] + ] + + assert equal ( + cartesian product [1, 2] [3, 4] [5, 6] + ) [ + [1, 3, 5], [1, 3, 6], [1, 4, 5], [1, 4, 6], [2, 3, 5], [2, 3, 6], [2, 4, 5], [2, 4, 6] + ] +}