implement count using mapreduce

This creates the same calling interface for `count` as for e.g. `sum`,
namely allowing the `dims` keyword.
The implementation is also shorter than before without sacrificing
performance.
`mapreduce` with `add_sum` may even yield performance benefits through
chunking, though this was not observed in simple tests.

NEWS.md

@@ -51,6 +51,8 @@ Standard library changes
 
 * Sets are now displayed less compactly in the REPL, as a column of elements, like vectors
   and dictionaries ([#33300]).
+* `count` now accepts the `dims` keyword.
+* new in-place `count!` function similar to `sum!`.
 
 #### Libdl
 

base/exports.jl

@@ -482,6 +482,7 @@ export
     any,
     firstindex,
     collect,
+    count!,
     count,
     delete!,
     deleteat!,

base/reduce.jl

@@ -846,18 +846,5 @@ julia> count([true, false, true, true])
 3
 ```
 """
-function count(pred, itr)
-    n = 0
-    for x in itr
-        n += pred(x)::Bool
-    end
-    return n
-end
-function count(pred, a::AbstractArray)
-    n = 0
-    for i in eachindex(a)
-        @inbounds n += pred(a[i])::Bool
-    end
-    return n
-end
 count(itr) = count(identity, itr)
+count(f, itr) = mapreduce(x->f(x)::Bool, add_sum, itr, init=0)

base/reducedim.jl

@@ -352,6 +352,63 @@ julia> reduce(max, a, dims=1)
 reduce(op, A::AbstractArray; kw...) = mapreduce(identity, op, A; kw...)
 
 ##### Specific reduction functions #####
+
+_bool(f::Function) = x->f(x)::Bool
+
+"""
+    count([f=identity,] A::AbstractArray; dims=:)
+
+Count the number of elements in `A` for which `f` returns `true` over the given
+dimensions.
+
+# Examples
+```jldoctest
+julia> A = [1 2; 3 4]
+2×2 Array{Int64,2}:
+ 1  2
+ 3  4
+
+julia> count(<=(2), A, dims=1)
+1×2 Array{Int64,2}:
+ 1  1
+
+julia> count(<=(2), A, dims=2)
+2×1 Array{Int64,2}:
+ 2
+ 0
+```
+"""
+count(A::AbstractArray; dims=:) = count(identity, A, dims=dims)
+count(f, A::AbstractArray; dims=:) = mapreduce(_bool(f), add_sum, A, dims=dims, init=0)
+
+"""
+    count!([f=identity,] r, A; init=true)
+
+Count the number of elements in `A` for which `f` returns `true` over the
+singleton dimensions of `r`, writing the result into `r` in-place.
+If `init` is `true`, values in `r` are initialized to zero.
+
+# Examples
+```jldoctest
+julia> A = [1 2; 3 4]
+2×2 Array{Int64,2}:
+ 1  2
+ 3  4
+
+julia> count!(<=(2), [1 1], A)
+1×2 Array{Int64,2}:
+ 1  1
+
+julia> count!(<=(2), [1; 1], A)
+2-element Array{Int64,1}:
+ 2
+ 0
+```
+"""
+count!(r::AbstractArray, A::AbstractArray; init::Bool=true) = count!(identity, r, A; init=init)
+count!(f, r::AbstractArray, A::AbstractArray; init::Bool=true) =
+    mapreducedim!(_bool(f), add_sum, initarray!(r, add_sum, init, A), A)
+
 """
     sum(A::AbstractArray; dims)