Skip to content

Commit a5b488b

Browse files
leerholeerho
authored
Merge pull request #760 from permutive-engineering/fix/quickselect-sorted-input
Choose the QuickSelect pivot by median of three
2 parents 325574e + 1a49f80 commit a5b488b

2 files changed

Lines changed: 111 additions & 22 deletions

File tree

‎src/main/java/org/apache/datasketches/common/QuickSelect.java‎

Lines changed: 39 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -107,20 +107,14 @@ public static long selectExcludingZeros(final long[] arr, final int nonZeros, fi
107107
* @return the next partition value. Ultimately, the desired pivot.
108108
*/
109109
private static int partition(final long[] arr, final int lo, final int hi) {
110-
int i = lo, j = hi + 1; //left and right scan indices
110+
medianOfThreeToLo(arr, lo, hi);
111+
int i = lo, j = hi; //left and right scan indices
111112
final long v = arr[lo]; //partitioning item value
112113
while (true) {
113114
//Scan right, scan left, check for scan complete, and exchange
114-
while (arr[ ++i] < v) {
115-
if (i == hi) {
116-
break;
117-
}
118-
}
119-
while (v < arr[ --j]) {
120-
if (j == lo) {
121-
break;
122-
}
123-
}
115+
//arr[hi] >= v and arr[lo] == v bound the scans
116+
while (arr[ ++i] < v) { }
117+
while (v < arr[ --j]) { }
124118
if (i >= j) {
125119
break;
126120
}
@@ -135,6 +129,21 @@ private static int partition(final long[] arr, final int lo, final int hi) {
135129
return j;
136130
}
137131

132+
//A pivot taken from a fixed position is quadratic on ordered input; a median of three is not.
133+
private static void medianOfThreeToLo(final long[] arr, final int lo, final int hi) {
134+
final int mid = lo + ((hi - lo) >>> 1);
135+
if (arr[mid] < arr[lo]) { swap(arr, mid, lo); }
136+
if (arr[hi] < arr[lo]) { swap(arr, hi, lo); }
137+
if (arr[hi] < arr[mid]) { swap(arr, hi, mid); }
138+
swap(arr, lo, mid);
139+
}
140+
141+
private static void swap(final long[] arr, final int i, final int j) {
142+
final long x = arr[i];
143+
arr[i] = arr[j];
144+
arr[j] = x;
145+
}
146+
138147
//For double arrays
139148

140149
/**
@@ -211,20 +220,14 @@ public static double selectExcludingZeros(final double[] arr, final int nonZeros
211220
* @return the next partition value. Ultimately, the desired pivot.
212221
*/
213222
private static int partition(final double[] arr, final int lo, final int hi) {
214-
int i = lo, j = hi + 1; //left and right scan indices
223+
medianOfThreeToLo(arr, lo, hi);
224+
int i = lo, j = hi; //left and right scan indices
215225
final double v = arr[lo]; //partitioning item value
216226
while (true) {
217227
//Scan right, scan left, check for scan complete, and exchange
218-
while (arr[ ++i] < v) {
219-
if (i == hi) {
220-
break;
221-
}
222-
}
223-
while (v < arr[ --j]) {
224-
if (j == lo) {
225-
break;
226-
}
227-
}
228+
//arr[hi] >= v and arr[lo] == v bound the scans
229+
while (arr[ ++i] < v) { }
230+
while (v < arr[ --j]) { }
228231
if (i >= j) {
229232
break;
230233
}
@@ -239,4 +242,18 @@ private static int partition(final double[] arr, final int lo, final int hi) {
239242
return j;
240243
}
241244

245+
private static void medianOfThreeToLo(final double[] arr, final int lo, final int hi) {
246+
final int mid = lo + ((hi - lo) >>> 1);
247+
if (arr[mid] < arr[lo]) { swap(arr, mid, lo); }
248+
if (arr[hi] < arr[lo]) { swap(arr, hi, lo); }
249+
if (arr[hi] < arr[mid]) { swap(arr, hi, mid); }
250+
swap(arr, lo, mid);
251+
}
252+
253+
private static void swap(final double[] arr, final int i, final int j) {
254+
final double x = arr[i];
255+
arr[i] = arr[j];
256+
arr[j] = x;
257+
}
258+
242259
}

‎src/test/java/org/apache/datasketches/thetacommon/QuickSelectTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,9 +53,29 @@ public void checkQuickSelect0Based() {
5353
final long retVal = select(arr, 0, len - 1, pivot);
5454
Assert.assertEquals(retVal, trueVal);
5555
}
56+
Assert.assertEquals(select(ascending(len), 0, len - 1, pivot), trueVal);
57+
Assert.assertEquals(select(descending(len), 0, len - 1, pivot), trueVal);
5658
}
5759
}
5860

61+
@Test
62+
public void checkQuickSelect0BasedAllEqual() {
63+
final int len = 64;
64+
for (int pivot = 0; pivot < len; pivot++ ) {
65+
final long[] arr = new long[len];
66+
java.util.Arrays.fill(arr, 7L);
67+
Assert.assertEquals(select(arr, 0, len - 1, pivot), 7L);
68+
}
69+
}
70+
71+
//Quadratic work still returns the right value, so only the time limit can fail here.
72+
@Test(timeOut = 10_000)
73+
public void checkOrderedInputIsNotQuadratic() {
74+
final int len = 1 << 19;
75+
Assert.assertEquals(select(ascending(len), 0, len - 1, len / 2), len / 2);
76+
Assert.assertEquals(select(descending(len), 0, len - 1, len / 2), len / 2);
77+
}
78+
5979
@Test
6080
public void checkQuickSelect1BasedExcludingZeros() {
6181
final int len = 64;
@@ -116,9 +136,28 @@ public void checkQuickSelectDbl0Based() {
116136
final double retVal = select(arr, 0, len - 1, pivot);
117137
Assert.assertEquals(retVal, trueVal, 0.0);
118138
}
139+
Assert.assertEquals(select(ascendingDbl(len), 0, len - 1, pivot), trueVal, 0.0);
140+
Assert.assertEquals(select(descendingDbl(len), 0, len - 1, pivot), trueVal, 0.0);
119141
}
120142
}
121143

144+
@Test
145+
public void checkQuickSelectDbl0BasedAllEqual() {
146+
final int len = 64;
147+
for (int pivot = 0; pivot < len; pivot++ ) {
148+
final double[] arr = new double[len];
149+
java.util.Arrays.fill(arr, 7.0);
150+
Assert.assertEquals(select(arr, 0, len - 1, pivot), 7.0, 0.0);
151+
}
152+
}
153+
154+
@Test(timeOut = 10_000)
155+
public void checkOrderedInputIsNotQuadraticDbl() {
156+
final int len = 1 << 19;
157+
Assert.assertEquals(select(ascendingDbl(len), 0, len - 1, len / 2), len / 2, 0.0);
158+
Assert.assertEquals(select(descendingDbl(len), 0, len - 1, len / 2), len / 2, 0.0);
159+
}
160+
122161
@Test
123162
public void checkQuickSelectDbl1BasedExcludingZeros() {
124163
final int len = 64;
@@ -164,6 +203,39 @@ public void checkQuickSelectDbl1BasedIncludingZeros() {
164203
}
165204

166205

206+
//select rearranges its input, so each call needs a fresh array
207+
private static long[] ascending(final int len) {
208+
final long[] a = new long[len];
209+
for (int i = 0; i < len; i++ ) {
210+
a[i] = i;
211+
}
212+
return a;
213+
}
214+
215+
private static long[] descending(final int len) {
216+
final long[] a = new long[len];
217+
for (int i = 0; i < len; i++ ) {
218+
a[i] = (len - 1) - i;
219+
}
220+
return a;
221+
}
222+
223+
private static double[] ascendingDbl(final int len) {
224+
final double[] a = new double[len];
225+
for (int i = 0; i < len; i++ ) {
226+
a[i] = i;
227+
}
228+
return a;
229+
}
230+
231+
private static double[] descendingDbl(final int len) {
232+
final double[] a = new double[len];
233+
for (int i = 0; i < len; i++ ) {
234+
a[i] = (len - 1) - i;
235+
}
236+
return a;
237+
}
238+
167239
/**
168240
* Rearrange the elements of an array in random order.
169241
* @param a long array

0 commit comments

Comments
 (0)