diff --git a/NEWS.md b/NEWS.md index b6ffe8c872af5..e944ab9409b20 100644 --- a/NEWS.md +++ b/NEWS.md @@ -48,6 +48,8 @@ Standard library changes * `view`, `@view`, and `@views` now work on `AbstractString`s, returning a `SubString` when appropriate ([#35879]). * All `AbstractUnitRange{<:Integer}`s now work with `SubString`, `view`, `@view` and `@views` on strings ([#35879]). * `sum`, `prod`, `maximum`, and `minimum` now support `init` keyword argument ([#36188], [#35839]). +* `unique(f, itr; seen=Set{T}())` now allows you to declare the container type used for + keeping track of values returned by `f` on elements of `itr` ([#36280]). #### LinearAlgebra * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]). diff --git a/base/set.jl b/base/set.jl index a3041ea927e3d..b98e65c844a1c 100644 --- a/base/set.jl +++ b/base/set.jl @@ -120,19 +120,25 @@ julia> unique(Real[1, 1.0, 2]) ``` """ function unique(itr) + if isa(IteratorEltype(itr), HasEltype) + T = eltype(itr) + out = Vector{T}() + seen = Set{T}() + for x in itr + if !in(x, seen) + push!(seen, x) + push!(out, x) + end + end + return out + end T = @default_eltype(itr) - out = Vector{T}() - seen = Set{T}() y = iterate(itr) - y === nothing && return out + y === nothing && return T[] x, i = y - if !isconcretetype(T) && IteratorEltype(itr) == EltypeUnknown() - S = typeof(x) - return _unique_from(itr, S[x], Set{S}((x,)), i) - end - push!(seen, x) - push!(out, x) - return unique_from(itr, out, seen, i) + S = typeof(x) + R = isconcretetype(T) ? T : S + return _unique_from(itr, R[x], Set{R}((x,)), i) end _unique_from(itr, out, seen, i) = unique_from(itr, out, seen, i) @@ -175,8 +181,18 @@ julia> unique(x -> x^2, [1, -1, 3, -3, 4]) 4 ``` """ -function unique(f, C) +function unique(f, C; seen::Union{Nothing,Set}=nothing) out = Vector{eltype(C)}() + if seen !== nothing + for x in C + y = f(x) + if y ∉ seen + push!(out, x) + push!(seen, y) + end + end + return out + end s = iterate(C) if s === nothing @@ -241,7 +257,7 @@ julia> unique!(iseven, [2, 3, 5, 7, 9]) 3 ``` """ -function unique!(f, A::AbstractVector) +function unique!(f, A::AbstractVector; seen::Union{Nothing,Set}=nothing) if length(A) <= 1 return A end @@ -249,7 +265,9 @@ function unique!(f, A::AbstractVector) i = firstindex(A) x = @inbounds A[i] y = f(x) - seen = Set{typeof(y)}() + if seen === nothing + seen = Set{typeof(y)}() + end push!(seen, y) return _unique!(f, A, seen, i, i+1) end diff --git a/test/sets.jl b/test/sets.jl index 5d8db78803cb8..3c213b5426c9d 100644 --- a/test/sets.jl +++ b/test/sets.jl @@ -384,27 +384,30 @@ end end @testset "unique" begin - u = unique([1, 1, 2]) + u = @inferred(unique([1, 1, 2])) @test in(1, u) @test in(2, u) @test length(u) == 2 - @test unique(iseven, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 8] - @test unique(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 1, 9] + @test @inferred(unique(iseven, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6])) == [5, 8] + @test @inferred(unique(x->x^2, Integer[3, -4, 5, 4])) == Integer[3, -4, 5] + @test @inferred(unique(iseven, Integer[3, -4, 5, 4]; seen=Set{Bool}())) == Integer[3, -4] + @test @inferred(unique(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6])) == [5, 1, 9] end @testset "issue 20105" begin @test @inferred(unique(x for x in 1:1)) == [1] @test unique(x for x in Any[1, 1.0])::Vector{Real} == [1] @test unique(x for x in Real[1, 1.0])::Vector{Real} == [1] - @test unique(Integer[1, 1, 2])::Vector{Integer} == [1, 2] + @test @inferred(unique(Integer[1, 1, 2]))::Vector{Integer} == [1, 2] + @test unique(x for x in []) isa Vector{Any} end @testset "unique!" begin u = [1,1,3,2,1] - unique!(u) + @inferred(unique!(u)) @test u == [1,3,2] - @test unique!([]) == [] - @test unique!(Float64[]) == Float64[] + @test @inferred(unique!([])) == [] + @test @inferred(unique!(Float64[])) == Float64[] u = [1,2,2,3,5,5] @test unique!(u) === u @test u == [1,2,3,5] @@ -434,8 +437,9 @@ end u = [1,2,5,1,3,2] @test unique!(x -> x ^ 2, [1, -1, 3, -3, 5, -5]) == [1, 3, 5] @test unique!(n -> n % 3, [5, 1, 8, 9, 3, 4, 10, 7, 2, 6]) == [5, 1, 9] - @test unique!(iseven, [2, 3, 5, 7, 9]) == [2, 3] - @test unique!(x -> x % 2 == 0 ? :even : :odd, [1, 2, 3, 4, 2, 2, 1]) == [1, 2] + @test @inferred(unique!(iseven, [2, 3, 5, 7, 9])) == [2, 3] + @test @inferred(unique!(x -> x % 2 == 0 ? :even : :odd, [1, 2, 3, 4, 2, 2, 1])) == [1, 2] + @test @inferred(unique!(x -> x % 2 == 0 ? :even : "odd", [1, 2, 3, 4, 2, 2, 1]; seen=Set{Union{Symbol,String}}())) == [1, 2] end @testset "allunique" begin