Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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;
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment