diff --git a/src/IterTools.jl b/src/IterTools.jl index 1211bb9..bebc7a3 100644 --- a/src/IterTools.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 6aa6372..ef3ca4f 100644 --- a/test/runtests.jl +++ b/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