Allow passing a function to extrema (#30323)

Currently `minimum` and `maximum` can accept a function argument, but
`extrema` cannot. This makes it consistent.

Author: ararslan

NEWS.md

@@ -4,6 +4,8 @@ Julia v1.2 Release Notes
 New language features
 ---------------------
 
+* The `extrema` function now accepts a function argument in the same manner as `minimum` and
+  `maximum` ([#30323]).
 
 Language changes
 ----------------
@@ -36,3 +38,4 @@ Deprecated or removed
 
 <!--- generated by NEWS-update.jl: -->
 [#29998]: https://github.com/JuliaLang/julia/issues/29998
+[#30323]: https://github.com/JuliaLang/julia/issues/30323

base/multidimensional.jl

@@ -1500,40 +1500,53 @@ julia> extrema(A, dims = (1,2))
  (9, 15)
 ```
 """
-extrema(A::AbstractArray; dims = :) = _extrema_dims(A, dims)
+extrema(A::AbstractArray; dims = :) = _extrema_dims(identity, A, dims)
 
-_extrema_dims(A::AbstractArray, ::Colon) = _extrema_itr(A)
+"""
+    extrema(f, A::AbstractArray; dims) -> Array{Tuple}
+
+Compute the minimum and maximum of `f` applied to each element in the given dimensions
+of `A`.
+
+!!! compat "Julia 1.2"
+    This method requires Julia 1.2 or later.
+"""
+extrema(f, A::AbstractArray; dims=:) = _extrema_dims(f, A, dims)
+
+_extrema_dims(f, A::AbstractArray, ::Colon) = _extrema_itr(f, A)
 
-function _extrema_dims(A::AbstractArray, dims)
+function _extrema_dims(f, A::AbstractArray, dims)
     sz = [size(A)...]
     for d in dims
         sz[d] = 1
     end
-    B = Array{Tuple{eltype(A),eltype(A)}}(undef, sz...)
-    return extrema!(B, A)
+    T = promote_op(f, eltype(A))
+    B = Array{Tuple{T,T}}(undef, sz...)
+    return extrema!(f, B, A)
 end
 
-@noinline function extrema!(B, A)
+@noinline function extrema!(f, B, A)
     @assert !has_offset_axes(B, A)
     sA = size(A)
     sB = size(B)
     for I in CartesianIndices(sB)
-        AI = A[I]
-        B[I] = (AI, AI)
+        fAI = f(A[I])
+        B[I] = (fAI, fAI)
     end
     Bmax = CartesianIndex(sB)
     @inbounds @simd for I in CartesianIndices(sA)
         J = min(Bmax,I)
         BJ = B[J]
-        AI = A[I]
-        if AI < BJ[1]
-            B[J] = (AI, BJ[2])
-        elseif AI > BJ[2]
-            B[J] = (BJ[1], AI)
+        fAI = f(A[I])
+        if fAI < BJ[1]
+            B[J] = (fAI, BJ[2])
+        elseif fAI > BJ[2]
+            B[J] = (BJ[1], fAI)
         end
     end
     return B
 end
+extrema!(B, A) = extrema!(identity, B, A)
 
 # Show for pairs() with Cartesian indices. Needs to be here rather than show.jl for bootstrap order
 function Base.showarg(io::IO, r::Iterators.Pairs{<:Integer, <:Any, <:Any, T}, toplevel) where T <: Union{AbstractVector, Tuple}

base/operators.jl

@@ -440,24 +440,43 @@ julia> extrema([9,pi,4.5])
 (3.141592653589793, 9.0)
 ```
 """
-extrema(itr) = _extrema_itr(itr)
+extrema(itr) = _extrema_itr(identity, itr)
 
-function _extrema_itr(itr)
+"""
+    extrema(f, itr) -> Tuple
+
+Compute both the minimum and maximum of `f` applied to each element in `itr` and return
+them as a 2-tuple. Only one pass is made over `itr`.
+
+!!! compat "Julia 1.2"
+    This method requires Julia 1.2 or later.
+
+# Examples
+```jldoctest
+julia> extrema(sin, 0:π)
+(0.0, 0.9092974268256817)
+```
+"""
+extrema(f, itr) = _extrema_itr(f, itr)
+
+function _extrema_itr(f, itr)
     y = iterate(itr)
     y === nothing && throw(ArgumentError("collection must be non-empty"))
     (v, s) = y
-    vmin = vmax = v
+    vmin = vmax = f(v)
     while true
         y = iterate(itr, s)
         y === nothing && break
         (x, s) = y
-        vmax = max(x, vmax)
-        vmin = min(x, vmin)
+        fx = f(x)
+        vmax = max(fx, vmax)
+        vmin = min(fx, vmin)
     end
     return (vmin, vmax)
 end
 
 extrema(x::Real) = (x, x)
+extrema(f, x::Real) = (y = f(x); (y, y))
 
 ## definitions providing basic traits of arithmetic operators ##
 

test/reduce.jl

@@ -180,10 +180,17 @@ prod2(itr) = invoke(prod, Tuple{Any}, itr)
 @test maximum(5) == 5
 @test minimum(5) == 5
 @test extrema(5) == (5, 5)
+@test extrema(abs2, 5) == (25, 25)
 
-@test maximum([4, 3, 5, 2]) == 5
-@test minimum([4, 3, 5, 2]) == 2
-@test extrema([4, 3, 5, 2]) == (2, 5)
+let x = [4,3,5,2]
+    @test maximum(x) == 5
+    @test minimum(x) == 2
+    @test extrema(x) == (2, 5)
+
+    @test maximum(abs2, x) == 25
+    @test minimum(abs2, x) == 4
+    @test extrema(abs2, x) == (4, 25)
+end
 
 @test isnan(maximum([NaN]))
 @test isnan(minimum([NaN]))
@@ -211,6 +218,7 @@ prod2(itr) = invoke(prod, Tuple{Any}, itr)
 
 @test maximum(abs2, 3:7) == 49
 @test minimum(abs2, 3:7) == 9
+@test extrema(abs2, 3:7) == (9, 49)
 
 @test maximum(Int16[1]) === Int16(1)
 @test maximum(Vector(Int16(1):Int16(100))) === Int16(100)
@@ -227,6 +235,7 @@ A = circshift(reshape(1:24,2,3,4), (0,1,1))
 @test size(extrema(A,dims=1)) == size(maximum(A,dims=1))
 @test size(extrema(A,dims=(1,2))) == size(maximum(A,dims=(1,2)))
 @test size(extrema(A,dims=(1,2,3))) == size(maximum(A,dims=(1,2,3)))
+@test extrema(x->div(x, 2), A, dims=(2,3)) == reshape([(0,11),(1,12)],2,1,1)
 
 # any & all