You can achieve your complexity goal by creating a histogram of the values:
public static IntStream maxValues(IntStream source, int limit) {
TreeMap<Integer,Integer> m=new TreeMap<>();
source.forEachOrdered(new IntConsumer() {
int size, min=Integer.MIN_VALUE;
public void accept(int value) {
if(value<min) return;
m.merge(value, 1, Integer::sum);
if(size<limit) size++;
else m.compute(min=m.firstKey(), (k,count)->count==1? null: count-1);
}
});
if(m.size()==limit)// no duplicates
return m.keySet().stream().mapToInt(Integer::valueOf);
return m.entrySet().stream().flatMapToInt(e->{
int value = e.getKey(), count = e.getValue();
return count==1? IntStream.of(value): IntStream.range(0, count).map(i->value);
});
}
It creates a map from int values to their corresponding number of occurrences but limits its contents to the desired number of values, hence, it’s operation has a O(log(M))
complexity (worst case, if no duplicates) and, since the operation is performed once for each value, it’s overall complexity is O(N×log(M))
as you wished.
You may test it with your original array as
int[] arr = {5, 3, 4, 2, 9, 1, 7, 8, 6};
maxValues(Arrays.stream(arr), 3).forEach(System.out::println);
but to test some corner cases, you may use an array containing duplicates like
int[] arr = {8, 5, 3, 4, 2, 2, 9, 1, 7, 9, 8, 6};
// note that the stream of three max elements contains one of the two eights
If you strive for maximum performance, replacing the boxing treemap with an adequate data structure using primitive data types may be feasible but that would be a minor performance optimization as this solution already solved the complexity problem.
By the way, this solution works for arbitrary streams, i.e. doesn’t need to know the value of N
.