what-the-bank/sources/com/google/common/collect/TopKSelector.java

170 lines
5.0 KiB
Java
Raw Permalink Normal View History

2024-07-27 18:17:47 +07:00
package com.google.common.collect;
import com.google.common.base.Preconditions;
import com.google.common.math.IntMath;
import java.math.RoundingMode;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
/* JADX INFO: Access modifiers changed from: package-private */
/* loaded from: classes2.dex */
public final class TopKSelector<T> {
private final T[] buffer;
private int bufferSize;
private final Comparator<? super T> comparator;
private final int k;
private T threshold;
public static <T extends Comparable<? super T>> TopKSelector<T> least(int i) {
return least(i, Ordering.natural());
}
public static <T> TopKSelector<T> least(int i, Comparator<? super T> comparator) {
return new TopKSelector<>(comparator, i);
}
public static <T extends Comparable<? super T>> TopKSelector<T> greatest(int i) {
return greatest(i, Ordering.natural());
}
public static <T> TopKSelector<T> greatest(int i, Comparator<? super T> comparator) {
return new TopKSelector<>(Ordering.from(comparator).reverse(), i);
}
private TopKSelector(Comparator<? super T> comparator, int i) {
this.comparator = (Comparator) Preconditions.checkNotNull(comparator, "comparator");
this.k = i;
Preconditions.checkArgument(i >= 0, "k (%s) must be >= 0", i);
Preconditions.checkArgument(i <= 1073741823, "k (%s) must be <= Integer.MAX_VALUE / 2", i);
this.buffer = (T[]) new Object[IntMath.checkedMultiply(i, 2)];
this.bufferSize = 0;
this.threshold = null;
}
public final void offer(T t) {
int i = this.k;
if (i == 0) {
return;
}
int i2 = this.bufferSize;
if (i2 == 0) {
this.buffer[0] = t;
this.threshold = t;
this.bufferSize = 1;
return;
}
if (i2 < i) {
T[] tArr = this.buffer;
this.bufferSize = i2 + 1;
tArr[i2] = t;
if (this.comparator.compare(t, this.threshold) > 0) {
this.threshold = t;
return;
}
return;
}
if (this.comparator.compare(t, this.threshold) < 0) {
T[] tArr2 = this.buffer;
int i3 = this.bufferSize;
int i4 = i3 + 1;
this.bufferSize = i4;
tArr2[i3] = t;
if (i4 == (this.k << 1)) {
trim();
}
}
}
private void trim() {
int i = (this.k << 1) - 1;
int log2 = IntMath.log2(i, RoundingMode.CEILING);
int i2 = 0;
int i3 = 0;
int i4 = 0;
while (true) {
if (i3 >= i) {
break;
}
int partition = partition(i3, i, ((i3 + i) + 1) >>> 1);
int i5 = this.k;
if (partition <= i5) {
if (partition >= i5) {
break;
}
i3 = Math.max(partition, i3 + 1);
i2 = partition;
} else {
i = partition - 1;
}
i4++;
if (i4 >= log2 * 3) {
Arrays.sort(this.buffer, i3, i, this.comparator);
break;
}
}
this.bufferSize = this.k;
this.threshold = this.buffer[i2];
while (true) {
i2++;
if (i2 >= this.k) {
return;
}
if (this.comparator.compare(this.buffer[i2], this.threshold) > 0) {
this.threshold = this.buffer[i2];
}
}
}
private int partition(int i, int i2, int i3) {
T[] tArr = this.buffer;
T t = tArr[i3];
tArr[i3] = tArr[i2];
int i4 = i;
while (i < i2) {
if (this.comparator.compare(this.buffer[i], t) < 0) {
swap(i4, i);
i4++;
}
i++;
}
T[] tArr2 = this.buffer;
tArr2[i2] = tArr2[i4];
tArr2[i4] = t;
return i4;
}
private void swap(int i, int i2) {
T[] tArr = this.buffer;
T t = tArr[i];
tArr[i] = tArr[i2];
tArr[i2] = t;
}
public final void offerAll(Iterable<? extends T> iterable) {
offerAll(iterable.iterator());
}
public final void offerAll(Iterator<? extends T> it) {
while (it.hasNext()) {
offer(it.next());
}
}
public final List<T> topK() {
Arrays.sort(this.buffer, 0, this.bufferSize, this.comparator);
int i = this.bufferSize;
int i2 = this.k;
if (i > i2) {
T[] tArr = this.buffer;
Arrays.fill(tArr, i2, tArr.length, (Object) null);
int i3 = this.k;
this.bufferSize = i3;
this.threshold = this.buffer[i3 - 1];
}
return Collections.unmodifiableList(Arrays.asList(Arrays.copyOf(this.buffer, this.bufferSize)));
}
}