Simplify combinations()

shubhamtagra · May 10, 2019 · d5dccb0 · d5dccb0
1 parent 0d9559f
commit d5dccb0
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/operator/scalar/ArrayCombinationsFunction.java b/presto-main/src/main/java/io/prestosql/operator/scalar/ArrayCombinationsFunction.java
@@ -70,12 +70,12 @@ public static Block combinations(
         int[] ids = new int[combinationCount * combinationLength];
         int idsPosition = 0;
 
-        int[] combination = firstCombination(combinationLength);
+        int[] combination = firstCombination(arrayLength, combinationLength);
         do {
             arraycopy(combination, 0, ids, idsPosition, combinationLength);
             idsPosition += combinationLength;
         }
-        while (nextCombination(combination, arrayLength));
+        while (nextCombination(combination, combinationLength));
         verify(idsPosition == ids.length, "idsPosition != ids.length, %s and %s respectively", idsPosition, ids.length);
 
         int[] offsets = new int[combinationCount + 1];
@@ -104,27 +104,23 @@ static int combinationCount(int arrayLength, int combinationLength)
         }
     }
 
-    private static int[] firstCombination(int combinationLength)
+    private static int[] firstCombination(int arrayLength, int combinationLength)
     {
-        int[] combination = new int[combinationLength];
+        int[] combination = new int[combinationLength + 1];
         setAll(combination, i -> i);
+        combination[combinationLength] = arrayLength; // sentinel
         return combination;
     }
 
-    private static boolean nextCombination(int[] combination, int arrayLength)
+    private static boolean nextCombination(int[] combination, int combinationLength)
     {
-        for (int i = 0; i < combination.length - 1; i++) {
+        for (int i = 0; i < combinationLength; i++) {
             if (combination[i] + 1 < combination[i + 1]) {
                 combination[i]++;
                 resetCombination(combination, i);
                 return true;
             }
         }
-        if (combination.length > 0 && combination[combination.length - 1] + 1 < arrayLength) {
-            combination[combination.length - 1]++;
-            resetCombination(combination, combination.length - 1);
-            return true;
-        }
         return false;
     }