Skip to content

Commit b1ceadc

Browse files
faster & type stable invperm and isperm for tuples < 16 elements (JuliaLang#35234)
* Make `invperm` much faster, and type stable for tuples Co-authored-by: Takafumi Arakaki <takafumi.a@gmail.com>
1 parent 4efb718 commit b1ceadc

File tree

2 files changed

+57
-2
lines changed

2 files changed

+57
-2
lines changed

base/combinatorics.jl

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,22 @@ end
3636

3737
# Basic functions for working with permutations
3838

39+
@inline function _foldoneto(op, acc, ::Val{N}) where N
40+
@assert N::Integer > 0
41+
if @generated
42+
quote
43+
acc_0 = acc
44+
Base.Cartesian.@nexprs $N i -> acc_{i} = op(acc_{i-1}, i)
45+
return $(Symbol(:acc_, N))
46+
end
47+
else
48+
for i in 1:N
49+
acc = op(acc, i)
50+
end
51+
return acc
52+
end
53+
end
54+
3955
"""
4056
isperm(v) -> Bool
4157
@@ -50,7 +66,9 @@ julia> isperm([1; 3])
5066
false
5167
```
5268
"""
53-
function isperm(A)
69+
isperm(A) = _isperm(A)
70+
71+
function _isperm(A)
5472
n = length(A)
5573
used = falses(n)
5674
for a in A
@@ -63,6 +81,18 @@ isperm(p::Tuple{}) = true
6381
isperm(p::Tuple{Int}) = p[1] == 1
6482
isperm(p::Tuple{Int,Int}) = ((p[1] == 1) & (p[2] == 2)) | ((p[1] == 2) & (p[2] == 1))
6583

84+
function isperm(P::Tuple)
85+
valn = Val(length(P))
86+
_foldoneto(true, valn) do b,i
87+
s = _foldoneto(false, valn) do s, j
88+
s || P[j]==i
89+
end
90+
b&s
91+
end
92+
end
93+
94+
isperm(P::Any16) = _isperm(P)
95+
6696
# swap columns i and j of a, in-place
6797
function swapcols!(a::AbstractMatrix, i, j)
6898
i == j && return
@@ -242,7 +272,21 @@ function invperm(p::Union{Tuple{},Tuple{Int},Tuple{Int,Int}})
242272
isperm(p) || throw(ArgumentError("argument is not a permutation"))
243273
p # in dimensions 0-2, every permutation is its own inverse
244274
end
245-
invperm(a::Tuple) = (invperm([a...])...,)
275+
276+
function invperm(P::Tuple)
277+
valn = Val(length(P))
278+
ntuple(valn) do i
279+
s = _foldoneto(nothing, valn) do s, j
280+
s !== nothing && return s
281+
P[j]==i && return j
282+
nothing
283+
end
284+
s === nothing && throw(ArgumentError("argument is not a permutation"))
285+
s
286+
end
287+
end
288+
289+
invperm(P::Any16) = Tuple(invperm(collect(P)))
246290

247291
#XXX This function should be moved to Combinatorics.jl but is currently used by Base.DSP.
248292
"""

test/combinatorics.jl

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,17 @@ end
4545
let ai = 2:-1:1
4646
@test invpermute!(permute!([1, 2], ai), ai) == [1, 2]
4747
end
48+
49+
# PR 35234
50+
for N in 3:1:20
51+
A=randcycle(N)
52+
T=Tuple(A)
53+
K=Tuple(A.-1)
54+
@test A[collect(invperm(T))] == 1:N
55+
@test_throws ArgumentError invperm(K)
56+
@test isperm(T) == true
57+
@test isperm(K) == false
58+
end
4859
end
4960

5061
@testset "factorial" begin

0 commit comments

Comments
 (0)