From acd0cb7b4d9ae026fb20c59337e0bd8e33ff1d96 Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <ifsmirnov@yandex.ru>
Date: Sun, 26 Mar 2017 01:22:45 +0300
Subject: [PATCH] randomTestSizes()

---
 array.h     |  6 +++++
 jngen.h     | 66 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 math.h      |  4 ++--
 testcases.h | 58 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 130 insertions(+), 4 deletions(-)

diff --git a/array.h b/array.h
index 6138295..9c59ebb 100644
--- a/array.h
+++ b/array.h
@@ -465,7 +465,13 @@ TArray<std::pair<T, U>> zip(const TArray<T>& lhs, const TArray<U>& rhs) {
     return result;
 }
 
+template<typename T, typename U>
+TArray<T> arrayCast(const TArray<U>& array) {
+    return TArray<T>(array.begin(), array.end());
+}
+
 } // namespace jngen
 
 using jngen::makeArray;
 using jngen::zip;
+using jngen::arrayCast;
diff --git a/jngen.h b/jngen.h
index eb48257..8d5650a 100644
--- a/jngen.h
+++ b/jngen.h
@@ -1984,10 +1984,16 @@ TArray<std::pair<T, U>> zip(const TArray<T>& lhs, const TArray<U>& rhs) {
     return result;
 }
 
+template<typename T, typename U>
+TArray<T> arrayCast(const TArray<U>& array) {
+    return TArray<T>(array.begin(), array.end());
+}
+
 } // namespace jngen
 
 using jngen::makeArray;
 using jngen::zip;
+using jngen::arrayCast;
 
 
 #include <algorithm>
@@ -2153,7 +2159,7 @@ class MathRandom {
         return Array(res.begin(), res.end());
     }
 
-    static Array64 partition(long long n, long long numParts) {
+    static Array64 partition(long long n, int numParts) {
         auto res = partitionNonEmpty(n + numParts, numParts);
         for (auto& x: res) {
             --x;
@@ -2167,7 +2173,7 @@ class MathRandom {
         return Array(res.begin(), res.end());
     }
 
-    static Array64 partitionNonEmpty(long long n, long long numParts) {
+    static Array64 partitionNonEmpty(long long n, int numParts) {
         ensure(numParts > 0);
         ensure(
             numParts <= n,
@@ -2796,6 +2802,7 @@ StringPair StringRandom::antiHash(
 using jngen::rnds;
 
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 
@@ -2833,11 +2840,66 @@ void setNextTestNumber(int testNo) {
     nextTestNo = testNo;
 }
 
+Array64 randomTestSizes(
+    long long totalSize,
+    int count,
+    long long minSize,
+    long long maxSize,
+    const Array64& predefined)
+{
+    for (auto x: predefined) {
+        totalSize -= x;
+    }
+    ensure(totalSize >= 0, "Sum of predefined test sizes exceeds total size");
+    ensure(count * minSize <= totalSize, "minSize is too large");
+    ensure(minSize <= maxSize);
+
+    auto partition = rndm.partition(totalSize - count * minSize, count)
+        .sort().reverse();
+    long long remaining = 0;
+
+    long long localMax = maxSize - minSize;
+    for (auto& x: partition) {
+        if (x > localMax) {
+            remaining += x - localMax;
+            x = localMax;
+        } else {
+            long long add = std::min(remaining, localMax - x);
+            x += add;
+            remaining -= add;
+        }
+
+        x += minSize;
+    }
+
+    ensure(remaining == 0, "maxSize is too small");
+
+    return (partition + predefined).shuffled();
+}
+
+Array randomTestSizes(
+    int totalSize,
+    int count,
+    int minSize,
+    int maxSize,
+    const Array& predefined)
+{
+    return arrayCast<int>(randomTestSizes(
+        static_cast<long long>(totalSize),
+        count,
+        static_cast<long long>(minSize),
+        static_cast<long long>(maxSize),
+        arrayCast<long long>(predefined)
+    ));
+}
+
 } // namespace jngen
 
 using jngen::startTest;
 using jngen::setNextTestNumber;
 
+using jngen::randomTestSizes;
+
 
 #include <cstring>
 #include <iostream>
diff --git a/math.h b/math.h
index 4859ada..8b5ba73 100644
--- a/math.h
+++ b/math.h
@@ -167,7 +167,7 @@ class MathRandom {
         return Array(res.begin(), res.end());
     }
 
-    static Array64 partition(long long n, long long numParts) {
+    static Array64 partition(long long n, int numParts) {
         auto res = partitionNonEmpty(n + numParts, numParts);
         for (auto& x: res) {
             --x;
@@ -181,7 +181,7 @@ class MathRandom {
         return Array(res.begin(), res.end());
     }
 
-    static Array64 partitionNonEmpty(long long n, long long numParts) {
+    static Array64 partitionNonEmpty(long long n, int numParts) {
         ensure(numParts > 0);
         ensure(
             numParts <= n,
diff --git a/testcases.h b/testcases.h
index 88a5770..b0a4a9e 100644
--- a/testcases.h
+++ b/testcases.h
@@ -1,7 +1,10 @@
 #pragma once
 
+#include "array.h"
 #include "common.h"
+#include "math.h"
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 
@@ -39,7 +42,62 @@ void setNextTestNumber(int testNo) {
     nextTestNo = testNo;
 }
 
+Array64 randomTestSizes(
+    long long totalSize,
+    int count,
+    long long minSize,
+    long long maxSize,
+    const Array64& predefined)
+{
+    for (auto x: predefined) {
+        totalSize -= x;
+    }
+    ensure(totalSize >= 0, "Sum of predefined test sizes exceeds total size");
+    ensure(count * minSize <= totalSize, "minSize is too large");
+    ensure(minSize <= maxSize);
+
+    auto partition = rndm.partition(totalSize - count * minSize, count)
+        .sort().reverse();
+    long long remaining = 0;
+
+    long long localMax = maxSize - minSize;
+    for (auto& x: partition) {
+        if (x > localMax) {
+            remaining += x - localMax;
+            x = localMax;
+        } else {
+            long long add = std::min(remaining, localMax - x);
+            x += add;
+            remaining -= add;
+        }
+
+        x += minSize;
+    }
+
+    ensure(remaining == 0, "maxSize is too small");
+
+    return (partition + predefined).shuffled();
+}
+
+Array randomTestSizes(
+    int totalSize,
+    int count,
+    int minSize,
+    int maxSize,
+    const Array& predefined)
+{
+    return arrayCast<int>(randomTestSizes(
+        static_cast<long long>(totalSize),
+        count,
+        static_cast<long long>(minSize),
+        static_cast<long long>(maxSize),
+        arrayCast<long long>(predefined)
+    ));
+}
+
 } // namespace jngen
 
 using jngen::startTest;
 using jngen::setNextTestNumber;
+
+using jngen::randomTestSizes;