Find neighbored segments/labels in Matrix

Here’s a simple implementation, let me know how it works :slight_smile:

function find_neighbors(A; nsegments=maximum(A))
    # These two should be compared to see which one performs better:
    flags = zeros(Bool, nsegments, nsegments)
    #flags = falses(nsegments, nsegments)  # Uses a BitArray instead of an array of Bool
    
    indices = CartesianIndices(A)
    neighbors = CartesianIndex(-1, -1):CartesianIndex(1, 1)
    for I in indices
        for N in neighbors
            if I+N in indices
                flags[A[I+N], A[I]] = true
            end
        end
    end

    # Unset diagonal elements since a segment is not neighbor of itself
    flags[diagind(flags)] .= false

    return Dict(seg => findall(col) for (seg, col) in enumerate(eachcol(flags)))
end
2 Likes