001/* *********************************************************************** *
002 * project: org.matsim.*
003 *                                                                         *
004 * *********************************************************************** *
005 *                                                                         *
006 * copyright       : (C) 2014 by the members listed in the COPYING,        *
007 *                   LICENSE and WARRANTY file.                            *
008 * email           : info at matsim dot org                                *
009 *                                                                         *
010 * *********************************************************************** *
011 *                                                                         *
012 *   This program is free software; you can redistribute it and/or modify  *
013 *   it under the terms of the GNU General Public License as published by  *
014 *   the Free Software Foundation; either version 2 of the License, or     *
015 *   (at your option) any later version.                                   *
016 *   See also COPYING, LICENSE and WARRANTY file                           *
017 *                                                                         *
018 * *********************************************************************** */
019
020package org.matsim.contrib.util;
021
022import java.util.Arrays;
023import java.util.List;
024import java.util.PriorityQueue;
025import java.util.function.ToDoubleFunction;
026import java.util.stream.Stream;
027
028/**
029 * Sorts from smallest to largest. If the opposite should be the case then add elements with their values negated:
030 * {@code PartialSort.add(element, -value)}. Works fine for small k (k << n); otherwise, one should consider a partial
031 * version of heapsort or quicksort.
032 * <p>
033 * More info: <a href="http://en.wikipedia.org/wiki/Partial_sorting">Partial sorting</a>
034 * 
035 * @param <T>
036 */
037public class PartialSort<T> {
038        public static <T> List<T> kSmallestElements(int k, Stream<T> elements, ToDoubleFunction<T> evaluator) {
039                PartialSort<T> nearestRequestSort = new PartialSort<T>(k);
040                nearestRequestSort.addAll(elements, evaluator::applyAsDouble);
041                return nearestRequestSort.kSmallestElements();
042        }
043
044        private static class ElementValuePair<T> implements Comparable<ElementValuePair<T>> {
045                private final T element;
046                private final double value;
047
048                public ElementValuePair(T element, double value) {
049                        this.element = element;
050                        this.value = value;
051                }
052
053                @Override
054                public int compareTo(ElementValuePair<T> o) {
055                        return -Double.compare(value, o.value);// reversed comparison (the smallest is the last in the queue)
056                }
057        }
058
059        private final int k;
060        private final PriorityQueue<ElementValuePair<T>> kSmallestElements;// descending order: from k-th to 1-st
061
062        public PartialSort(int k) {
063                this.k = k;
064                kSmallestElements = new PriorityQueue<>(k);
065        }
066
067        public void add(T element, double value) {
068                if (kSmallestElements.size() < k) {
069                        kSmallestElements.add(new ElementValuePair<>(element, value));
070                } else if (Double.compare(value, kSmallestElements.peek().value) < 0) {
071                        kSmallestElements.poll();
072                        kSmallestElements.add(new ElementValuePair<>(element, value));
073                }
074        }
075
076        public void addAll(Stream<T> elements, ToDoubleFunction<T> evaluator) {
077                elements.forEach(e -> this.add(e, evaluator.applyAsDouble(e)));
078        }
079
080        /**
081         * Gets k smallest elements (side effect: they are removed from the queue -- the queue gets empty).
082         * 
083         * @return list containing k smallest elements sorted ascending: from the smallest to the k-th smallest
084         */
085        public List<T> kSmallestElements() {
086                @SuppressWarnings("unchecked")
087                T[] array = (T[])new Object[kSmallestElements.size()];
088                for (int i = array.length - 1; i >= 0; i--) {
089                        array[i] = kSmallestElements.poll().element;
090                }
091                return Arrays.asList(array);
092        }
093}