add sliding_window_maxima

src/IterTools.jl

@@ -35,7 +35,8 @@ export
     fieldvalues,
     interleaveby,
     cache,
-    zip_longest
+    zip_longest,
+    sliding_window_maxima
 
 function has_length(it)
     it_size = IteratorSize(it)
@@ -1235,4 +1236,168 @@ t = ('e', 'x')
 """
 zip_longest(its...; default=nothing) = ZipLongest(Tuple(_Padded.(its, default)))
 
+module ValidatedPositiveInt
+    export validated_positive_int
+    @noinline function throw_err_not_positive()
+        throw(ArgumentError("not positive"))
+    end
+    function validated_positive_int(m::Int)
+        if m < 1
+            @noinline throw_err_not_positive()
+        end
+        m
+    end
+end
+
+module DeleteFromEnd
+    export delete_from_end!
+    @noinline function cannot_delete_more_than_length()
+        throw(ArgumentError("the collection does not have that many elements"))
+    end
+    function delete_from_end!(c::Vector, n::Int)
+        l = length(c)
+        if l < n
+            @noinline cannot_delete_more_than_length()
+        end
+        for _ ∈ 1:n
+            pop!(c)
+        end
+        c
+    end
+end
+
+module SlidingWindowMaximumIterators
+    export SlidingWindowMaximumIterator
+    using ..ValidatedPositiveInt, ..DeleteFromEnd
+    struct SlidingWindowMaximumIterator{Iterator, Ord <: Base.Order.Ordering}
+        window_size::Int
+        iterator::Iterator
+        order::Ord
+        function SlidingWindowMaximumIterator(window_size::Int, iterator, order::Base.Order.Ordering)
+            s = validated_positive_int(window_size)
+            new{typeof(iterator), typeof(order)}(s, iterator, order)
+        end
+    end
+    function Base.IteratorSize(::Type{<:SlidingWindowMaximumIterator})
+        Base.SizeUnknown()
+    end
+    function Base.IteratorEltype(::Type{<:SlidingWindowMaximumIterator{Iterator}}) where {Iterator}
+        Base.IteratorEltype(Iterator)
+    end
+    function Base.eltype(::Type{<:SlidingWindowMaximumIterator{Iterator}}) where {Iterator}
+        eltype(Iterator)
+    end
+    function get_window_size(iterator::SlidingWindowMaximumIterator)
+        iterator.window_size
+    end
+    function delete_all_lesser_from_end!(window_queue::Vector, order::Base.Order.Ordering, elem)
+        pop_count = 0
+        len = length(window_queue)
+        while (pop_count < len) && Base.Order.lt(order, window_queue[end - pop_count][1], elem)
+            pop_count = pop_count + 1
+        end
+        delete_from_end!(window_queue, pop_count)
+    end
+    # `window_queue` is logically a double-ended queue data structure: only mutating it
+    # with `pop!`, `popfirst` and `push!`.
+    function _iterate(iterator::SlidingWindowMaximumIterator, state::Tuple{Tuple{(Vector{Tuple{T, Int}} where {T}), Int, Any}})
+        (window_queue, counter, inner_iterator_state_initial) = state[1]
+        counter = counter::Int
+        iter = iterate(iterator.iterator, inner_iterator_state_initial)
+        if iter === nothing
+            return iter
+        end
+        (elem, inner_iterator_state) = iter
+        order = iterator.order
+        delete_all_lesser_from_end!(window_queue, order, elem)
+        if (!isempty(window_queue)) && (get_window_size(iterator) ≤ counter - window_queue[1][2]::Int)
+            popfirst!(window_queue)
+        end
+        push!(window_queue, (elem, counter))
+        next_state = (window_queue, counter + 1, inner_iterator_state)
+        (window_queue[1][1], (next_state,))
+    end
+    function _iterate(iterator::SlidingWindowMaximumIterator, ::Tuple{})
+        inner_iterator = iterator.iterator
+        iter_initial = iterate(inner_iterator)
+        if iter_initial === nothing
+            return iter_initial
+        end
+        (elem_initial, inner_iterator_state) = iter_initial
+        window_size = get_window_size(iterator)
+        counter = 0
+        window_queue = [(elem_initial, counter)]
+        order = iterator.order
+        for _ ∈ 2:window_size
+            iter = iterate(inner_iterator, inner_iterator_state)
+            if iter === nothing
+                return iter
+            end
+            (elem, inner_iterator_state) = iter
+            delete_all_lesser_from_end!(window_queue, order, elem)
+            counter = counter + 1
+            push!(window_queue, (elem, counter))
+        end
+        state = (window_queue, counter + 1, inner_iterator_state)
+        (window_queue[1][1], (state,))
+    end
+    function Base.iterate(iterator::SlidingWindowMaximumIterator, state = ())
+        _iterate(iterator, state)
+    end
+end
+
+"""
+    sliding_window_maxima(window_size::Number, iterator, [order::Base.Order.Ordering])
+
+An iterator. Each element is the maximum of a sliding window of size `window_size`, with
+the original elements being taken from the other iterator, `iterator`. `window_size`
+must be convertible to `Int`.
+
+```jldoctest
+julia> v = Float32[1, 7, 7, 4, 9]
+5-element Vector{Float32}:
+ 1.0
+ 7.0
+ 7.0
+ 4.0
+ 9.0
+
+julia> collect(sliding_window_maxima(1, v)) == v
+true
+
+julia> collect(sliding_window_maxima(2, v))
+4-element Vector{Float32}:
+ 7.0
+ 7.0
+ 7.0
+ 9.0
+
+julia> collect(sliding_window_maxima(3, v))
+3-element Vector{Float32}:
+ 7.0
+ 7.0
+ 9.0
+```
+
+The optional argument `order` determines how the maximum is computed.
+
+```jldoctest
+julia> collect(sliding_window_maxima(3, 1:5))
+3-element Vector{Int64}:
+ 3
+ 4
+ 5
+
+julia> collect(sliding_window_maxima(3, 1:5, Base.Order.Reverse))
+3-element Vector{Int64}:
+ 1
+ 2
+ 3
+```
+"""
+function sliding_window_maxima(window_size::Number, iterator, order::Base.Order.Ordering = Base.Order.Forward)
+    s = Int(window_size)
+    SlidingWindowMaximumIterators.SlidingWindowMaximumIterator(s, iterator, order)
+end
+
 end # module IterTools

test/runtests.jl

@@ -612,5 +612,18 @@ include("testing_macros.jl")
         @test collect(it_mixed) == [(1,10),(2,9),(3,8),(4,' '),(5, ' ')]
         @test eltype(it_mixed) == Tuple{Union{Missing, Int}, Union{Char, Int}}
     end
+
+    @testset "sliding_window_maxima" begin
+        vec = Float32[1, 2, 3, 3, 4, 5, 4]
+        iter = @inferred sliding_window_maxima(2, vec)
+        @test (@inferred collect(iter)) isa Vector{Float32}
+        @test collect(sliding_window_maxima(1, vec))::Vector{Float32} == vec
+        @test collect(sliding_window_maxima(2, vec))::Vector{Float32} == [2, 3, 3, 4, 5, 5]
+        @test collect(sliding_window_maxima(3, vec))::Vector{Float32} == [3, 3, 4, 5, 5]
+        @test collect(sliding_window_maxima(3, vec, Base.Order.Reverse))::Vector{Float32} == [1, 2, 3, 3, 4]
+        @test collect(sliding_window_maxima(100, vec))::Vector{Float32} == []
+        @test collect(sliding_window_maxima(1, Float32[]))::Vector{Float32} == []
+        @test_throws ArgumentError sliding_window_maxima(-1, vec)
+    end
 end
 end