01mf02 / jaq

A jq clone focussed on correctness, speed, and simplicity
MIT License
2.87k stars 77 forks source link

combinations/1 #35

Open pkoppstein opened 2 years ago

pkoppstein commented 2 years ago

Here is a jaq-defined definition of combinations/1 with the same semantics as jq's combinations/1. It seems to me that the helper functions, namely decimal2base/1 and combinations/2, are both independently worthy of inclusion in the jaq library (and could in fact also be included in the jq library in the sense that they behave in the same way when run using jq), so I have not folded them into the def of combinations/1.

# Input: a positive integer
# Output: an array representing the number in base b, with the least significant digit first.
def decimal2base(b):
  b as $b
  | [recurse(if . > 0 then ./$b|floor else empty end) | . % $b]
  | if length > 1 then .[:-1] else . end;

# Enumerate all the ways to select m elements from range(0;n) with replacement.
# The output is a stream of arrays of length m.
def combinations(n; m):
 n as $n
 | m as $m
 | [1, []]  # state: [i, combination]
 | while( (.[1] | length) <= m;
          .[0] | [. + 1, decimal2base($n)] )
 | .[1]
 | [range(0; $m - length) | 0]  + reverse;

def combinations(n):
  combinations(length; n) as $c
  | [ .[$c[]] ] ; 

Example:

["a", "b"] | combinations(3)

01mf02 commented 2 years ago

decimal2base looks nice already; can you make it so that the input can be negative as well?

pkoppstein commented 1 year ago

Now that jaq supports recursion as used in jq's def combination:, jaq can simply include jq's defs for combinations/0 and combinations/1.