Created
May 17, 2018 16:12
-
-
Save kweimann/1ea76a35e1667aa2a74c072e163efaf2 to your computer and use it in GitHub Desktop.
K smallest / largest elements in O(N * log(K)) complexity (N - size of the collection)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
object implicits { | |
implicit class RichTraversable[A](traversable: Traversable[A]) { | |
// collect k smallest elements in O(n*log(k)) time (n - length of the list, k - size of the window) | |
def minK(k: Int)(implicit cmp: Ordering[A]): List[A] = { | |
if (k < 1) return List.empty | |
traversable | |
.foldLeft(mutable.PriorityQueue.empty) { | |
// window not filled yet so append current element | |
case (window, x) if window.size < k => | |
window.enqueue(x) | |
window | |
// current element is bigger than any element in the window so discard it | |
case (window, x) if cmp.compare(x, window.head) >= 0 => | |
window | |
// discard biggest element from the window and append current element | |
case (window, x) => | |
window.dequeue() | |
window.enqueue(x) | |
window | |
} | |
.toList | |
.sorted | |
} | |
// collect k largest elements in O(n*log(k)) time (n - length of the list, k - size of the window) | |
def maxK(k: Int)(implicit cmp: Ordering[A]): List[A] = { | |
// semantically same as the min method but uses reverse ordering | |
minK(k)((x: A, y: A) => cmp.compare(y, x)) | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment