Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 71 additions & 0 deletions stdlib/SparseArrays/src/sparsematrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1369,6 +1369,77 @@ function sparse_sortedlinearindices!(I::Vector{Ti}, V::Vector, m::Int, n::Int) w
return SparseMatrixCSC(m, n, colptr, I, V)
end

# findfirst/next/prev/last
function _idxfirstnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedfirst(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxnextnz(A, searchk)
end

function _idxlastnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedlast(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxprevnz(A, searchk)
end

function _idxnextnz(A::SparseMatrixCSC, idx::Integer)
nnza = nnz(A)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx <= nnza
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx += 1
end
return zero(idx), z
end

function _idxprevnz(A::SparseMatrixCSC, idx::Integer)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx > 0
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx -= 1
end
return zero(idx), z
end

function _idx_to_cartesian(A::SparseMatrixCSC, idx::Integer)
rowval = rowvals(A)
i = rowval[idx]
j = searchsortedlast(A.colptr, idx, 1, size(A, 2), Base.Order.Forward)
return CartesianIndex(i, j)
end

function Base.findnext(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findnext, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxfirstnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxnextnz(A, idx + 1)
end
return nothing
end

function Base.findprev(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findprev, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxlastnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxprevnz(A, idx - 1)
end
return nothing
end

"""
sprand([rng],[type],m,[n],p::AbstractFloat,[rfn])

Expand Down
23 changes: 19 additions & 4 deletions stdlib/SparseArrays/test/sparse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2245,16 +2245,20 @@ end
@test findprev(!iszero, x,i) == findprev(!iszero, x_sp,i)
end

y = [0 0 0 0 0;
y = [7 0 0 0 0;
1 0 1 0 0;
1 0 0 0 1;
1 7 0 7 1;
0 0 1 0 0;
1 0 1 1 0]
y_sp = sparse(y)
1 0 1 1 0.0]
y_sp = [x == 7 ? -0.0 : x for x in sparse(y)]
y = Array(y_sp)
@test isequal(y_sp[1,1], -0.0)

for i in keys(y)
@test findnext(!iszero, y,i) == findnext(!iszero, y_sp,i)
@test findprev(!iszero, y,i) == findprev(!iszero, y_sp,i)
@test findnext(iszero, y,i) == findnext(iszero, y_sp,i)
@test findprev(iszero, y,i) == findprev(iszero, y_sp,i)
end

z_sp = sparsevec(Dict(1=>1, 5=>1, 8=>0, 10=>1))
Expand All @@ -2264,6 +2268,17 @@ end
@test findnext(!iszero, z,i) == findnext(!iszero, z_sp,i)
@test findprev(!iszero, z,i) == findprev(!iszero, z_sp,i)
end

w = [ "a" ""; "" "b"]
w_sp = sparse(w)

for i in keys(w)
@test findnext(!isequal(""), w,i) == findnext(!isequal(""), w_sp,i)
@test findprev(!isequal(""), w,i) == findprev(!isequal(""), w_sp,i)
@test findnext(isequal(""), w,i) == findnext(isequal(""), w_sp,i)
@test findprev(isequal(""), w,i) == findprev(isequal(""), w_sp,i)
end

end

# #20711
Expand Down