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}