non-staged sub2ind/ind2sub improvement

JuliaLang · Feb 26, 2015 · 2b60d66 · 2b60d66
1 parent 8b75161
commit 2b60d66
Showing 1 changed file with 13 additions and 50 deletions.
diff --git a/base/abstractarray.jl b/base/abstractarray.jl
@@ -1030,28 +1030,11 @@ function repmat(a::AbstractVector, m::Int)
     return b
 end
 
-sub2ind(dims) = 1
-sub2ind(dims, i::Integer) = int(i)
-sub2ind(dims, i::Integer, j::Integer) = sub2ind(dims, int(i), int(j))
-sub2ind(dims, i::Int, j::Int) = (j-1)*dims[1] + i
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer) = sub2ind(dims, int(i0),int(i1),int(i2))
-sub2ind(dims, i0::Int, i1::Int, i2::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*(i2-1))
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer, i3::Integer) =
-    sub2ind(dims, int(i0),int(i1),int(i2),int(i3))
-sub2ind(dims, i0::Int, i1::Int, i2::Int, i3::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*((i2-1) + dims[3]*(i3-1)))
-
-function sub2ind(dims, I::Integer...)
-    ndims = length(dims)
-    index = int(I[1])
-    stride = 1
-    for k=2:ndims
-        stride = stride * dims[k-1]
-        index += (int(I[k])-1) * stride
-    end
-    return index
-end
+sub2ind(dims::()) = 1
+sub2ind(dims::(Integer,), i1::Integer) = i1
+sub2ind(dims::(Integer,Integer), i1::Integer, i2::Integer) = i1+dims[1]*(i2-1)
+sub2ind(dims::(Integer,Integer,Integer...), i1::Integer, i2::Integer, I::Integer...) =
+    i1 + dims[1]*(sub2ind(tail(dims),i2,I...)-1)
 
 function sub2ind{T<:Integer}(dims::Array{T}, sub::Array{T})
     ndims = length(dims)
@@ -1064,37 +1047,18 @@ function sub2ind{T<:Integer}(dims::Array{T}, sub::Array{T})
     return ind
 end
 
-sub2ind{T<:Integer}(dims, I::AbstractVector{T}...) =
+sub2ind{T<:Integer}(dims::(Integer...,), I::AbstractVector{T}...) =
     [ sub2ind(dims, map(X->X[i], I)...)::Int for i=1:length(I[1]) ]
 
-function ind2sub(dims::(Integer,Integer...), ind::Int)
-    ndims = length(dims)
-    stride = dims[1]
-    for i=2:ndims-1
-        stride *= dims[i]
-    end
-
-    sub = ()
-    for i=(ndims-1):-1:1
-        rest = rem(ind-1, stride) + 1
-        sub = tuple(div(ind - rest, stride) + 1, sub...)
-        ind = rest
-        stride = div(stride, dims[i])
-    end
-    return tuple(ind, sub...)
-end
 
-ind2sub(dims::(Integer...), ind::Integer) = ind2sub(dims, int(ind))
-ind2sub(dims::(), ind::Integer) = ind==1 ? () : throw(BoundsError())
-ind2sub(dims::(Integer,), ind::Int) = (ind,)
-ind2sub(dims::(Integer,Integer), ind::Int) =
-    (rem(ind-1,dims[1])+1, div(ind-1,dims[1])+1)
-ind2sub(dims::(Integer,Integer,Integer), ind::Int) =
-    (rem(ind-1,dims[1])+1, div(rem(ind-1,dims[1]*dims[2]), dims[1])+1,
-     div(rem(ind-1,dims[1]*dims[2]*dims[3]), dims[1]*dims[2])+1)
-ind2sub(a::AbstractArray, ind::Integer) = ind2sub(size(a), Int(ind))
+ind2sub(dims::(), ind::Integer) = ()
+ind2sub(dims::(Integer,), ind::Integer) = (ind,)
+ind2sub(dims::(Integer,Integer), ind::Integer) = (rem(ind-1,dims[1])+1,div(ind-1,dims[1])+1)
+ind2sub(dims::(Integer,Integer,Integer...), ind::Integer) =
+    tuple(rem(ind-1,dims[1])+1, ind2sub(tail(dims),div(ind-1,dims[1])+1)...)
+ind2sub(a::AbstractArray, ind::Integer) = ind2sub(size(a), ind)
 
-function ind2sub{T<:Integer}(dims::(Integer,Integer...), ind::AbstractVector{T})
+function ind2sub{T<:Integer}(dims::(Integer...), ind::AbstractVector{T})
     n = length(dims)
     l = length(ind)
     t = ntuple(n, x->Array(Int, l))
@@ -1499,4 +1463,3 @@ function randsubseq!(S::AbstractArray, A::AbstractArray, p::Real)
 end
 
 randsubseq{T}(A::AbstractArray{T}, p::Real) = randsubseq!(T[], A, p)
-