Merge pull request #925 from HarlanH/master

findfirst(A) returns index of first non-zero element, or 0
JuliaLang · Jun 19, 2012 · 7816b4a · 7816b4a
2 parents e1bd5a0 + b020893
commit 7816b4a
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 0 deletions.
diff --git a/base/array.jl b/base/array.jl
@@ -894,6 +894,67 @@ function nnz(a::StridedArray)
     return n
 end
 
+# returns the index of the first non-zero element, or 0 if all zeros
+function findfirst{T}(A::StridedArray{T})
+    z = zero(T)
+    for i = 1:length(A)
+        if A[i] != z
+            return i
+        end
+    end
+    return 0
+end
+
+# returns the index of the first matching element
+function findfirst{T}(A::StridedArray{T}, v::T)
+    for i = 1:length(A)
+        if A[i] == v
+            return i
+        end
+    end
+    return 0
+end
+
+# returns the index of the first element for which the function returns true
+function findfirst{T}(A::StridedArray{T}, testf::Function)
+    for i = 1:length(A)
+        if testf(A[i])
+            return i
+        end
+    end
+    return 0
+end
+
+
+function find{T}(A::StridedArray{T}, v::T)
+    # use a dynamic-length array to store the indexes, then copy to a non-padded
+    # array for the return
+    tmpI = Array(Int, 0)
+    for i = 1:length(A)
+        if A[i] == v
+            push(tmpI, i)
+        end
+    end
+    I = Array(Int, length(tmpI))
+    copy_to(I, tmpI)
+    I
+end
+
+function find{T}(A::StridedArray{T}, testf::Function)
+    # use a dynamic-length array to store the indexes, then copy to a non-padded
+    # array for the return
+    tmpI = Array(Int, 0)
+    for i = 1:length(A)
+        if testf(A[i])
+            push(tmpI, i)
+        end
+    end
+    I = Array(Int, length(tmpI))
+    copy_to(I, tmpI)
+    I
+end
+
+
 function find(A::StridedArray)
     nnzA = nnz(A)
     I = Array(Int, nnzA)

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -88,6 +88,19 @@ Y = [2, 1, 4, 3]
 @assert X[Y[end],1] == 5
 @assert X[end,Y[end]] == 11
 
+## find, findfirst ##
+a = [0,1,2,3,0,1,2,3]
+@assert find(a) == [2,3,4,6,7,8]
+@assert find(a.==2) == [3,7]
+@assert find(a,2) == [3,7]
+@assert find(a,isodd) == [2,4,6,8]
+@assert findfirst(a) == 2
+@assert findfirst(a.==0) == 1
+@assert findfirst(a.==5) == 0
+@assert findfirst([1,2,4,1,2,3,4], 3) == 6
+@assert findfirst([2,4,6,3,9,2,0], isodd) == 4
+@assert findfirst([2,4,6,2,0], isodd) == 0
+
 
 ## findn ##