From b8d860cebb772bd9d58ddd268966d6f10bacbbe0 Mon Sep 17 00:00:00 2001 From: Simeon Schaub Date: Tue, 8 Sep 2020 13:12:36 +0200 Subject: [PATCH 1/2] add init argument to count Co-authored-by: Rafael Fourquet --- NEWS.md | 1 + base/bitarray.jl | 8 ++++---- base/reduce.jl | 32 +++++++++++++++++++------------- base/reducedim.jl | 11 +++++++---- test/bitarray.jl | 2 ++ test/reduce.jl | 6 ++++++ test/reducedim.jl | 12 ++++++++++++ 7 files changed, 51 insertions(+), 21 deletions(-) diff --git a/NEWS.md b/NEWS.md index c37415f0206e3..3001bd998fc65 100644 --- a/NEWS.md +++ b/NEWS.md @@ -102,6 +102,7 @@ New library features inserting or consuming the first dimension depending on the ratio of `sizeof(T)` and `sizeof(S)`. * New `append!(vector, collections...)` and `prepend!(vector, collections...)` methods accept multiple collections to be appended or prepended ([#36227]). +* `count` now accepts an optional `init` argument to control the accumulation type ([#37461]). Standard library changes ------------------------ diff --git a/base/bitarray.jl b/base/bitarray.jl index 51eff74a3f50f..e51f617852b15 100644 --- a/base/bitarray.jl +++ b/base/bitarray.jl @@ -1386,15 +1386,15 @@ circshift!(B::BitVector, i::Integer) = circshift!(B, B, i) ## count & find ## -function bitcount(Bc::Vector{UInt64}) - n = 0 +function bitcount(Bc::Vector{UInt64}; init::T=0) where {T} + n::T = init @inbounds for i = 1:length(Bc) - n += count_ones(Bc[i]) + n += count_ones(Bc[i]) % T end return n end -count(B::BitArray) = bitcount(B.chunks) +count(B::BitArray; init=0) = bitcount(B.chunks; init) function unsafe_bitfindnext(Bc::Vector{UInt64}, start::Int) chunk_start = _div64(start-1)+1 diff --git a/base/reduce.jl b/base/reduce.jl index 4a6f27a5d947a..0b049c4429e7e 100644 --- a/base/reduce.jl +++ b/base/reduce.jl @@ -938,12 +938,15 @@ end _bool(f) = x->f(x)::Bool """ - count(p, itr) -> Integer - count(itr) -> Integer + count([f=identity,] itr; init=0) -> Integer -Count the number of elements in `itr` for which predicate `p` returns `true`. -If `p` is omitted, counts the number of `true` elements in `itr` (which -should be a collection of boolean values). +Count the number of elements in `itr` for which the function `f` returns `true`. +If `f` is omitted, counts the number of `true` elements in `itr` (which +should be a collection of boolean values). `init` optionally specifies the value +to start counting from and therefore also determines the output type. + +!!! compat "Julia 1.6" + `init` keyword was added in Julia 1.6. # Examples ```jldoctest @@ -952,32 +955,35 @@ julia> count(i->(4<=i<=6), [2,3,4,5,6]) julia> count([true, false, true, true]) 3 + +julia> count(>(3), 1:7, init=0x00) +0x04 ``` """ -count(itr) = count(identity, itr) +count(itr; init=0) = count(identity, itr; init) -count(f, itr) = _simple_count(f, itr) +count(f, itr; init=0) = _simple_count(f, itr, init) -function _simple_count(pred, itr) - n = 0 +function _simple_count(pred, itr, init::T) where {T} + n::T = init for x in itr n += pred(x)::Bool end return n end -function count(::typeof(identity), x::Array{Bool}) - n = 0 +function _simple_count(::typeof(identity), x::Array{Bool}, init::T=0) where {T} + n::T = init chunks = length(x) รท sizeof(UInt) mask = 0x0101010101010101 % UInt GC.@preserve x begin ptr = Ptr{UInt}(pointer(x)) for i in 1:chunks - n += count_ones(unsafe_load(ptr, i) & mask) + n = (n + count_ones(unsafe_load(ptr, i) & mask)) % T end end for i in sizeof(UInt)*chunks+1:length(x) - n += x[i] + n += x[i] % T end return n end diff --git a/base/reducedim.jl b/base/reducedim.jl index c889392ececd7..c83c2c71cb3e6 100644 --- a/base/reducedim.jl +++ b/base/reducedim.jl @@ -369,6 +369,9 @@ dimensions. !!! compat "Julia 1.5" `dims` keyword was added in Julia 1.5. +!!! compat "Julia 1.6" + `init` keyword was added in Julia 1.6. + # Examples ```jldoctest julia> A = [1 2; 3 4] @@ -386,11 +389,11 @@ julia> count(<=(2), A, dims=2) 0 ``` """ -count(A::AbstractArrayOrBroadcasted; dims=:) = count(identity, A, dims=dims) -count(f, A::AbstractArrayOrBroadcasted; dims=:) = _count(f, A, dims) +count(A::AbstractArrayOrBroadcasted; dims=:, init=0) = count(identity, A; dims, init) +count(f, A::AbstractArrayOrBroadcasted; dims=:, init=0) = _count(f, A, dims, init) -_count(f, A::AbstractArrayOrBroadcasted, dims::Colon) = _simple_count(f, A) -_count(f, A::AbstractArrayOrBroadcasted, dims) = mapreduce(_bool(f), add_sum, A, dims=dims, init=0) +_count(f, A::AbstractArrayOrBroadcasted, dims::Colon, init) = _simple_count(f, A, init) +_count(f, A::AbstractArrayOrBroadcasted, dims, init) = mapreduce(_bool(f), add_sum, A; dims, init) """ count!([f=identity,] r, A) diff --git a/test/bitarray.jl b/test/bitarray.jl index 06cc99e525747..ebc99e5fe9568 100644 --- a/test/bitarray.jl +++ b/test/bitarray.jl @@ -1219,6 +1219,8 @@ timesofar("datamove") @check_bit_operation findall(falses(t)) ret_type @check_bit_operation findall(bitrand(t)) ret_type end + + @test count(trues(2, 2), init=0x03) === 0x07 end timesofar("find") diff --git a/test/reduce.jl b/test/reduce.jl index 4688709b099c9..4f9fc33403282 100644 --- a/test/reduce.jl +++ b/test/reduce.jl @@ -520,6 +520,12 @@ struct NonFunctionIsZero end @test count(NonFunctionIsZero(), [0]) == 1 @test count(NonFunctionIsZero(), [1]) == 0 +@test count(Iterators.repeated(true, 3), init=0x04) === 0x07 +@test count(!=(2), Iterators.take(1:7, 3), init=Int32(0)) === Int32(2) +@test count(identity, [true, false], init=Int8(5)) === Int8(6) +@test count(!, [true false; false true], dims=:, init=Int16(0)) === Int16(2) +@test isequal(count(identity, [true false; false true], dims=2, init=UInt(4)), reshape(UInt[5, 5], 2, 1)) + ## cumsum, cummin, cummax z = rand(10^6) diff --git a/test/reducedim.jl b/test/reducedim.jl index ae5a253b22001..b6229634b6006 100644 --- a/test/reducedim.jl +++ b/test/reducedim.jl @@ -77,6 +77,15 @@ safe_minabs(A::Array{T}, region) where {T} = safe_mapslices(minimum, abs.(A), re @test @inferred(maximum(abs, Areduc, dims=region)) โ‰ˆ safe_maxabs(Areduc, region) @test @inferred(minimum(abs, Areduc, dims=region)) โ‰ˆ safe_minabs(Areduc, region) @test @inferred(count(!, Breduc, dims=region)) โ‰ˆ safe_count(.!Breduc, region) + + @test isequal( + @inferred(count(Breduc, dims=region, init=0x02)), + safe_count(Breduc, region) .% UInt8 .+ 0x02, + ) + @test isequal( + @inferred(count(!, Breduc, dims=region, init=Int16(0))), + safe_count(.!Breduc, region) .% Int16, + ) end # Combining dims and init @@ -446,3 +455,6 @@ end @test_throws TypeError count([1], dims=1) @test_throws TypeError count!([1], [1]) end + +@test @inferred(count(false:true, dims=:, init=0x0004)) === 0x0005 +@test @inferred(count(isodd, reshape(1:9, 3, 3), dims=:, init=Int128(0))) === Int128(5) From 4a1a8c888b82da35168c0703ff898b583f311baf Mon Sep 17 00:00:00 2001 From: Simeon Schaub Date: Sat, 17 Oct 2020 11:59:20 +0200 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Rafael Fourquet --- base/bitarray.jl | 2 +- base/reduce.jl | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/base/bitarray.jl b/base/bitarray.jl index e51f617852b15..94bb94e5a4d03 100644 --- a/base/bitarray.jl +++ b/base/bitarray.jl @@ -1389,7 +1389,7 @@ circshift!(B::BitVector, i::Integer) = circshift!(B, B, i) function bitcount(Bc::Vector{UInt64}; init::T=0) where {T} n::T = init @inbounds for i = 1:length(Bc) - n += count_ones(Bc[i]) % T + n = (n + count_ones(Bc[i])) % T end return n end diff --git a/base/reduce.jl b/base/reduce.jl index 0b049c4429e7e..8ea928669ab9e 100644 --- a/base/reduce.jl +++ b/base/reduce.jl @@ -941,7 +941,7 @@ _bool(f) = x->f(x)::Bool count([f=identity,] itr; init=0) -> Integer Count the number of elements in `itr` for which the function `f` returns `true`. -If `f` is omitted, counts the number of `true` elements in `itr` (which +If `f` is omitted, count the number of `true` elements in `itr` (which should be a collection of boolean values). `init` optionally specifies the value to start counting from and therefore also determines the output type. @@ -956,8 +956,8 @@ julia> count(i->(4<=i<=6), [2,3,4,5,6]) julia> count([true, false, true, true]) 3 -julia> count(>(3), 1:7, init=0x00) -0x04 +julia> count(>(3), 1:7, init=0x03) +0x07 ``` """ count(itr; init=0) = count(identity, itr; init) @@ -983,7 +983,7 @@ function _simple_count(::typeof(identity), x::Array{Bool}, init::T=0) where {T} end end for i in sizeof(UInt)*chunks+1:length(x) - n += x[i] % T + n = (n + x[i]) % T end return n end