Quickselect.cs

using System.Collections.Generic;

namespace System.Linq.Enumerable
{
	public static class Quickselect
	{
		/// <summary>
		/// Selects <b>count</b> <i>smallest</i> elements according to a key. Result is returned in no particular order.
		/// </summary>
		/// <typeparam name="T">Type of elements in collection</typeparam>
		/// <typeparam name="K">Type of key for comparison</typeparam>
		/// <param name="seq">Source collection from which the values are being selected</param>
		/// <param name="count">Maximum number of elements to be returned</param>
		/// <param name="keySelector">Function to map element to key for comparison</param>
		/// <param name="comparer">Comparer to compare key values</param>
		/// <returns>New list of maximum <b>count</b> <i>smallest</i> elements</returns>
		public static List<T> QuickselectSmallest<T, K>(this IEnumerable<T> seq, int count, Func<T, K> keySelector, IComparer<K> comparer = null)
		{
			return QuickselectInternal(seq, count, true, keySelector, comparer);
		}

		/// <summary>
		/// Selects <b>count</b> <i>largest</i> elements according to a key. Result is returned in no particular order.
		/// </summary>
		/// <typeparam name="T">Type of elements in collection</typeparam>
		/// <typeparam name="K">Type of key for comparison</typeparam>
		/// <param name="seq">Source collection from which the values are being selected</param>
		/// <param name="count">Maximum number of elements to be returned</param>
		/// <param name="keySelector">Function to map element to key for comparison</param>
		/// <param name="comparer">Comparer to compare key values</param>
		/// <returns>New list of maximum <b>count</b> <i>largest</i> elements</returns>
		public static List<T> QuickselectLargest<T, K>(this IEnumerable<T> seq, int count, Func<T, K> keySelector, IComparer<K> comparer = null)
		{
			return QuickselectInternal(seq, count, false, keySelector, comparer);
		}

		private static List<T> QuickselectInternal<T, K>(this IEnumerable<T> seq, int count, bool lookingForSmallest, Func<T, K> keySelector, IComparer<K> comparer)
		{
			if (count < 0)
				throw new ArgumentOutOfRangeException("count", "Count is smaller than 0.");
			if (count == 0)
				return new List<T>(0);

			var partiallySortedArray = seq.ToList();
			if (partiallySortedArray.Count <= count)
				return partiallySortedArray;

			if (comparer == null) comparer = Comparer<K>.Default;
			return QuickSelect(partiallySortedArray, count - 1, lookingForSmallest, keySelector, comparer).Take(count).ToList();
		}

		private static List<T> QuickSelect<T, K>(List<T> partiallySortedArray, int n, bool lookingForSmallest, Func<T, K> keySelector, IComparer<K> comparer)
		{
			var startIndex = 0;
			var endIndex = partiallySortedArray.Count - 1;
			var pivotIndex = n;
			var r = new Random();
			while (endIndex > startIndex)
			{
				pivotIndex = Partition(partiallySortedArray, startIndex, endIndex, pivotIndex, lookingForSmallest, keySelector, comparer);
				if (pivotIndex == n)
					break;
				if (pivotIndex > n)
					endIndex = pivotIndex - 1;
				else
					startIndex = pivotIndex + 1;
				pivotIndex = r.Next(startIndex, endIndex);
			}
			return partiallySortedArray;
		}

		private static int Partition<T, K>(this List<T> collection, int startIndex, int endIndex, int pivotIndex, bool lookingForSmallest, Func<T, K> keySelector, IComparer<K> comparer)
		{
			var pivotValue = keySelector(collection[pivotIndex]);
			collection.Swap(pivotIndex, endIndex);
			for (var i = startIndex; i < endIndex; i++)
			{
				var compareValue = comparer.Compare(keySelector(collection[i]), pivotValue);
				if ((lookingForSmallest ? compareValue : -compareValue) > 0)
					continue;
				collection.Swap(i, startIndex);
				startIndex++;
			}
			collection.Swap(endIndex, startIndex);
			return startIndex;
		}

		private static void Swap<T>(this List<T> list, int index1, int index2)
		{
			if (index1 == index2)
				return;

			var temp = list[index1];
			list[index1] = list[index2];
			list[index2] = temp;
		}
	}
}