by and lt arguments to nlargest/nsmallest

Author: jlumpe

src/heaps.jl

@@ -129,37 +129,45 @@ function nextreme(ord::Base.Ordering, n::Int, arr::AbstractVector{T}) where T
 end
 
 """
-    nlargest(n, arr)
+    nlargest(n, arr; kw...)
 
 Return the `n` largest elements of the array `arr`.
 
 Equivalent to:
-    sort(arr, order = Base.Reverse)[1:min(n, end)]
+    sort(arr, kw..., rev=true)[1:min(n, end)]
 
 Note that if `arr` contains floats and is free of NaN values,
-then the following alternative may be used to achieve 2x performance.
+then the following alternative may be used to achieve 2x performance:
+
     DataStructures.nextreme(DataStructures.FasterReverse(), n, arr)
+
 This faster version is equivalent to:
+
     sort(arr, lt = >)[1:min(n, end)]
 """
-function nlargest(n::Int, arr::AbstractVector)
-    return nextreme(Base.Reverse, n, arr)
+function nlargest(n::Int, arr::AbstractVector; lt=isless, by=identity)
+    order = Base.ReverseOrdering(Base.ord(lt, by, nothing))
+    return nextreme(order, n, arr)
 end
 
 """
-    nsmallest(n, arr)
+    nsmallest(n, arr; kw...)
 
 Return the `n` smallest elements of the array `arr`.
 
 Equivalent to:
-    sort(arr, order = Base.Forward)[1:min(n, end)]
+    sort(arr; kw...)[1:min(n, end)]
 
 Note that if `arr` contains floats and is free of NaN values,
-then the following alternative may be used to achieve 2x performance.
+then the following alternative may be used to achieve 2x performance:
+
     DataStructures.nextreme(DataStructures.FasterForward(), n, arr)
+
 This faster version is equivalent to:
+
     sort(arr, lt = <)[1:min(n, end)]
 """
-function nsmallest(n::Int, arr::AbstractVector)
-    return nextreme(Base.Forward, n, arr)
+function nsmallest(n::Int, arr::AbstractVector; lt=isless, by=identity)
+    order = Base.ord(lt, by, nothing)
+    return nextreme(order, n, arr)
 end

test/test_binheap.jl

@@ -208,9 +208,17 @@
             102,-56,-17,-41,25,-30,-84,26,-84,48,49,-5,-38,28,
             114,-54,96,-55,67,74,127,-61,124,11,-7,93,-51,110,
             -106,-84,-90,-18,-12,-116,21,115,50]
+        square = x -> x^2
+
         for n = -1:length(ss) + 1
-            @test sort(ss, lt = >)[1:min(n, end)] == nlargest(n, ss)
-            @test sort(ss, lt = <)[1:min(n, end)] == nsmallest(n, ss)
+            r = 1:min(n, length(ss))
+
+            @test nlargest(n, ss) == sort(ss, rev=true)[r]
+            @test nsmallest(n, ss) == sort(ss)[r]
+
+            @test nlargest(n, ss, by=square) == sort(ss, by=square, rev=true)[r]
+            @test nsmallest(n, ss, by=square) == sort(ss, by=square)[r]
+
             @test nlargest(n, ss) == DataStructures.nextreme(DataStructures.FasterReverse(), n, ss)
             @test nsmallest(n, ss) == DataStructures.nextreme(DataStructures.FasterForward(), n, ss)
         end