#!/usr/bin/env python3 from collections import Counter from functools import lru_cache import argparse import math import random N = 5000 S = 10**6 from lib import * def product_from_to(a, b): if b-a < 5: res = 1 for c in range(a,b): res *= c return res c = (a+b) // 2 return product_from_to(a,c) * product_from_to(c,b) def binomial(n, k): if k < 0 or n < k: return 0 k = min(k, n-k) return product_from_to(n-k+1, n+1) // product_from_to(1, k+1) # We say that `partitions(n,s,w)` is the number of ways to split a sequence of # `s` characters into `n` words, each of length at most `w`. # Relevant OEIS sequence: https://oeis.org/A273975 # Straightforward but inefficient recursive implementation: @lru_cache(10**6) def partitions_rec(n, w, s): if s < n or s > n*w: return 0 if n == 0: return 1 return sum([partitions_rec(n-1, w, s-k) for k in range(1,w+1)]) # More efficient recursive formula, can be derived using generating functions: def partitions(n, w, s): if s < n or s > n*w: return 0 if n == 0: return 1 s = min(s, n*(w+1)-s) res = 0 for i in range(min(n, (s-n)//w) + 1): j = s-n-w*i res += binomial(n,i) * (1 if i%2==0 else -1) * binomial(n+j-1,j) return res # Approximate time complexity of computing `partitions(n,w,s)`: def partitions_cost(n, w, s): s = min(s, n*(w+1)-s) return min(n, (s-n)//w) * n # For large inputs, use a rough approximate of the logarithm instead: ys = [0, 0.676, 0.786, 0.849, 0.894, 0.924, 0.955, 0.973, 0.986, 0.994, 1] def partitions_log(n, w, s): s = min(s, n*(w+1)-s) t = (s-n)/(n*w-n) i = int(t*20) f = ys[i] + (20*t-i)*(ys[i+1]-ys[i]) if i < 10 else 1 return n * math.log(w) * f # Generate a random partition. The distribution is uniform for small inputs. # For large inputs, the distribution is not uniform due to the approximations # but it should be somewhat close. def gen_partition(n, w, s): if w == 1: return [1] * n ks = [] while n > 0: kmin = max(1, s-w*(n-1)) kmax = min(w, s-(n-1)) p = (s-n) / ((w-1)*n) if partitions_cost(n,w,s) <= 1000: total = partitions(n,w,s) weights = [partitions(n-1,w,s-k)/total for k in range(kmin,kmax+1)] else: total_log = partitions_log(n,w,s) weights = [math.exp(partitions_log(n-1,w,s-k) - total_log) for k in range(kmin, kmax+1)] k = random.choices(range(kmin,kmax+1), weights=weights)[0] ks.append(k) n -= 1 s -= k random.shuffle(ks) return ks # Check if a valid test case for given `(n,w,s)` exists. def is_valid(n, w, s): if s < n or s > n*w: return False ks = sum([[k] * (26**k) for k in range(1,4)], [])[:n] return max(ks) <= w and sum(ks) <= s # Generate a test case by first generating a random partition for `(n,w,s)` and # trying to generate a corresponding list of words. This may fail if there are # too many short strings (there may be at most 26 words of length 1, at most # 676 words of length of 2, etc.). In this case, redistribute the characters # between the strings until a valid test case exists: def gen_case(n, w, s): try: ks = gen_partition(n,w,s) return random_from_lengths(ks) except ValueError: c = Counter(ks) for k in range(5): while c[k] > 26**k: kmax = max(c) assert k+1 < kmax, 'Impossible to create a test case with these parameters.' c = +(c - Counter([k,kmax]) + Counter([k+1,kmax-1])) words = [] alphabet = [chr(ord('a')+k) for k in range(26)] for k,m in c.items(): new_words = set() while len(new_words) < m: new_words.add(''.join(random.choices(alphabet, k=k))) words += list(new_words) return sorted(words) # Parse arguments and initialize random seed: parser = argparse.ArgumentParser(description='Generate a random test case.') parser.add_argument('-n', dest='n', type=int, required=True, help='the number of words') parser.add_argument('-w', dest='w', type=int, help='the width of the test case') parser.add_argument('-m', dest='m', type=int, help='the maximal word length') parser.add_argument('-s', dest='s', type=int, help='the total length of all words') parser.add_argument('-seed', dest='seed', type=int, required=True, help='the random seed') args = parser.parse_args() n, w, m, s, seed = args.n, args.w, args.m, args.s, args.seed random.seed(seed) # Randomly choose values for omitted arguments: ks = sum([[k] * (26**k) for k in range(1,4)], [])[:n] if m is None: mmin = max(ks) if s is None else max(max(ks), math.ceil(s/n)) mmax = N if w is None else w t = random.randint(1,5) if mmax >= t*mmin: mmax //= t assert mmin <= mmax, 'Impossible to create a test case with these parameters.' m = random.randint(mmin, mmax) if s is None: smin = sum(sum([[k] * (26**k) for k in range(1,4)], [])[:n]) smax = min(n*m,S) assert smin <= smax, 'Impossible to create a test case with these parameters.' s = int(round(random.triangular(smin,smax,smin))) if w is None: wmin = 2*m if 2*m <= N else m wmax = min(s//2 if s//2 >= wmin else 10*wmin, N) w = random.randint(wmin, wmax) assert is_valid(n, m, s), 'Impossible to create a test case with these parameters.' # Generate and print the test case: words = gen_case(n, m, s) print(n, w) print('\n'.join(words))