Fastest way to partition array, given a condition

To conclude: I tested, hopefully correctly, the alternative partition of quick_sort, but for my specific case the original implementation I posted above seems to be faster (I renamed it partition!, and here I apply it to a data structure which is the actual one I’m using):

julia> @benchmark partition!(s,f) setup=setup evals=1
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  209.000 ns …  17.891 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     382.000 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   412.933 ns ± 261.169 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

      ▁▄███▆▂
  ▂▂▃▅███████▇▅▄▃▃▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▁▂▁▁▂▁▁▁▁▁▁▁▁▁▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂ ▃
  209 ns           Histogram: frequency by time         1.53 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

julia> @benchmark partition_quicksort!(s,f) setup=setup evals=1
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  243.000 ns …   9.613 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     435.000 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   471.085 ns ± 226.825 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

       ▃▆██▇▅▂
  ▂▃▃▅█████████▆▅▄▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂ ▃
  243 ns           Histogram: frequency by time         1.56 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

Thus, I’m sticking to the original implementation. The field! option is very interesting, and I will be using it in an alternative implementation that allows a different data structure and might be faster in specific scenarios. (if anyone is curious, these are for optimizations of this package)

code

using StaticArrays, BenchmarkTools

function partition!(v::AbstractVector, by)
  iswap = 1
  @inbounds for i in eachindex(v)
    if by(v[i])
      if iswap != i
        v[iswap], v[i] = v[i], v[iswap]
      end
      iswap += 1
    end
  end
  return iswap - 1
end

function partition_quicksort!(v::AbstractVector, by)
    i, j = extrema(eachindex(v))
    i -= 1
    j += 1
    @inbounds while true
        i += 1; j -= 1
        while i <= length(v) && by(v[i]) 
          i += 1  
        end
        while j > 0 && !by(v[j])
          j -= 1 
        end
        i >= j && break
        v[i], v[j] = v[j], v[i]
    end
    return j
end

struct S
  i::Int64
  x::Float64
  v::SVector{3,Float32}
end

function check()
  for i in 1:100
    s = [ S(rand(1:20), rand(), rand(SVector{3,Float64})) for i in 1:rand(1:20) ]  
    s1 = deepcopy(s)
    s2 = deepcopy(s)
    cutoff = -0.1 + 1.1*rand()
    n1 = partition!(s1, el -> el.x <= cutoff)
    n2 = partition_quicksort!(s2, el -> el.x <= cutoff)
    if n1 != n2 || !(sum(x->x.x,s1) ≈ sum(x->x.x,s2))
      error()
    end
  end
end

check()

setup=(s = [ S(rand(1:20), rand(), rand(SVector{3,Float64})) for i in 1:20 ]; f = el -> el.x <= -0.1 + 1.1*rand() )

@benchmark partition!(s,f) setup=setup evals=1

@benchmark partition_quicksort!(s,f) setup=setup evals=1

1 Like

This might be obvious, but it hasn’t been mentioned yet. There is a partialsort function in Base.

julia> v = [4, 6, 1, 5, 7, 3, 2]
7-element Vector{Int64}:
 4
 6
 1
 5
 7
 3
 2

julia> partialsort(v, 1:3)
3-element view(::Vector{Int64}, 1:3) with eltype Int64:
 1
 2
 3

julia> partialsort(v, 1:3, rev=true)
3-element view(::Vector{Int64}, 1:3) with eltype Int64:
 7
 6
 5

partialsort! exists, but it may be possible to speed it up:

julia> @benchmark partialsort_cutoff!(a, b) setup=(c=1000;a=shuffle(1:c); b=rand(1:c))
BechmarkTools.Trial: 10000 samples with 10 evaluations.                               
 Range (min … max):  1.330 μs … 183.190 μs  ┊ GC (min … max): 0.00% … 0.00%           
 Time  (median):     1.720 μs               ┊ GC (median):    0.00%                   
 Time  (mean ± σ):   1.939 μs ±   2.531 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%           
                                                                                      
     ▂▄█▃▂                                                                            
  ▃▅██████▃▃▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂                       
  1.33 μs         Histogram: frequency by time         5.4 μs <                       
                                                                                      
 Memory estimate: 0 bytes, allocs estimate: 0.                                        
                                                                                      
julia> @benchmark partialsort!(a, 1:b) setup=(c=1000;a=shuffle(1:c); b=rand(1:c))     
BechmarkTools.Trial: 10000 samples with 122 evaluations.                              
 Range (min … max):  835.246 ns … 36.746 μs  ┊ GC (min … max): 0.00% … 0.00%          
 Time  (median):       3.048 μs              ┊ GC (median):    0.00%                  
 Time  (mean ± σ):     3.229 μs ±  2.142 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%          
                                                                                      
  ▆▇ ▇▄ ▅▄▇█   ▄▂▆▅▁▆▅▄▆▁                                                             
  █████▇█████▆████████████▇▄▄▃▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▄                      
  835 ns          Histogram: frequency by time         10.5 μs <                      
                                                                                      
 Memory estimate: 0 bytes, allocs estimate: 0.                                        

That said, the partialsort_cutoff! function isn’t very generic and assumes 1-based indexing, so there’d be some work needed before a PR :slight_smile:

2 Likes

Seems to me that partialsort! uses one valid index of the array as input, like of what is needed for the pivot in Base.Sort.partition!. That is not exactly the same thing I need here, which is sorting by an external condition. I have the impresion that to use that I would need first to know which is the element of the list which is closer to the boundary of the partition, but would take almost the same time as partitioning itself.

No, as I understand it partialsort in Base takes either an index or a range, such that the index is at the correct position or that the values at the indices given by the range are sorted (which is what you want, if you give it 1:i as I understand it).

Unless I am confused by something else, no, I don’t want that. I want that the elements get split into two blocks, the ones with a property smaller than something and the ones with a property greater than something. I don’t know in advance the number of elements of each group, neither if there is any element in either group. So I cannot provide an index nor a range as the input.

(as pointed above, what I want is to “partition” the array, I didn’t know that this was the name of the operation, maybe I should (I will) change the name of the thread, to avoid confusion).

2 Likes

Ah, I see! My apologies. You’d have to supply the index the element you’d like to have as your pivot would be at in the sorted array, which is why partialsort doesn’t work for you.

1 Like