Skip to content

Commit

Permalink
base/sort: add sort! for multidimensional arrays (#28902)
Browse files Browse the repository at this point in the history
  • Loading branch information
stev47 authored and StefanKarpinski committed Sep 11, 2018
1 parent 70c7298 commit 7142712
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 0 deletions.
49 changes: 49 additions & 0 deletions base/sort.jl
Original file line number Diff line number Diff line change
Expand Up @@ -938,6 +938,8 @@ Sort a multidimensional array `A` along the given dimension.
See [`sort!`](@ref) for a description of possible
keyword arguments.
To sort slices of an array, refer to [`sortslices`](@ref).
# Examples
```jldoctest
julia> A = [4 3; 1 2]
Expand Down Expand Up @@ -987,6 +989,53 @@ end
Av
end

"""
sort!(A; dims::Integer, alg::Algorithm=defalg(v), lt=isless, by=identity, rev::Bool=false, order::Ordering=Forward)
Sort the multidimensional array `A` along dimension `dims`.
See [`sort!`](@ref) for a description of possible keyword arguments.
To sort slices of an array, refer to [`sortslices`](@ref).
# Examples
```jldoctest
julia> A = [4 3; 1 2]
2×2 Array{Int64,2}:
4 3
1 2
julia> sort!(A, dims = 1); A
2×2 Array{Int64,2}:
1 2
4 3
julia> sort!(A, dims = 2); A
2×2 Array{Int64,2}:
1 2
3 4
```
"""
function sort!(A::AbstractArray;
dims::Integer,
alg::Algorithm=defalg(A),
lt=isless,
by=identity,
rev::Union{Bool,Nothing}=nothing,
order::Ordering=Forward)
ordr = ord(lt, by, rev, order)
nd = ndims(A)
k = dims

1 <= k <= nd || throw(ArgumentError("dimension out of range"))

remdims = ntuple(i -> i == k ? 1 : size(A, i), nd)
for idx in CartesianIndices(remdims)
Av = view(A, ntuple(i -> i == k ? Colon() : idx[i], nd)...)
sort!(Av, alg, ordr)
end
A
end

## fast clever sorting for floats ##

module Float
Expand Down
5 changes: 5 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1199,11 +1199,13 @@ end
@test issorted(as[:,1])
@test issorted(as[:,2])
@test issorted(as[:,3])
@test sort!(copy(a), dims=1) == as

as = sort(a, dims=2)
@test issorted(as[1,:])
@test issorted(as[2,:])
@test issorted(as[3,:])
@test sort!(copy(a), dims=2) == as

local b = rand(21,21,2)

Expand All @@ -1212,15 +1214,18 @@ end
@test issorted(bs[:,i,1])
@test issorted(bs[:,i,2])
end
@test sort!(copy(b), dims=1) == bs

bs = sort(b, dims=2)
for i in 1:21
@test issorted(bs[i,:,1])
@test issorted(bs[i,:,2])
end
@test sort!(copy(b), dims=2) == bs

bs = sort(b, dims=3)
@test all(bs[:,:,1] .<= bs[:,:,2])
@test sort!(copy(b), dims=3) == bs
end

@testset "higher dimensional sortslices" begin
Expand Down

0 comments on commit 7142712

Please sign in to comment.