Efficiently remove diagonal from Array

I’m trying to remove the diagonal from a two-dimensional square Array. I have come up with the code below but it’s not very efficient (e.g., on a 1000 x 1000 matrix, it allocates 75MiB and takes about 65ms on my machine). Is there a more efficient way?

function remove_diagonal(x)
    mat = Array{Float64}(undef, size(x, 1), size(x, 1) - 1)
    for i = 1:size(x, 1)
        inds = setdiff(1:size(x, 1), i)
        mat[i, :] .= x[i, inds]
    end
    return mat
end

Like this?

julia> a = [1 2 3; 4 5 6; 7 8 9]
3×3 Array{Int64,2}:
 1  2  3
 4  5  6
 7  8  9

julia> using LinearAlgebra

julia> Diagonal(a)
3×3 Diagonal{Int64,Array{Int64,1}}:
 1  ⋅  ⋅
 ⋅  5  ⋅
 ⋅  ⋅  9

Or this:

julia> tril(a, -1) + triu(a, 1)
3×3 Array{Int64,2}:
 0  2  3
 4  0  6
 7  8  0

Or this mutating, faster, version:

julia> for ij in CartesianIndices(a)
       if ij.I[1] == ij.I[2]
       a[ij] = 0
       end
       end

julia> a
3×3 Array{Int64,2}:
 0  2  3
 4  0  6
 7  8  0

I think the objective is not to zero out the diagonal, but to remove it.

I think that removing an element from each column instead of one from each row would give better memory access patterns (when iterating over columns), but I am not sure if that part of the spec can be changed. See

https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-column-major-1

@yakir12, thanks! But how do I then efficiently get from

0  2  3
4  0  6
7  8  0

to

2 3
4 6
7 8

?

You could spell it out more explicitly:


function remove_diagonal2(x)
    mat = Array{Float64}(undef, size(x, 1), size(x, 1) - 1)
    for i = 1:size(mat, 1)
        for j ∈ 1:size(mat, 2)
            if i > j
                mat[i, j] = x[i, j]
            else
                mat[i, j] = x[i, j+1]
            end
        end    
    end
    return mat
end

On my machine this gives me (using @btime remove_diagonal2($xl) where xl = rand(1000,1000):

# remove_diagonal
60.985 ms (19002 allocations: 75.87 MiB)
# remove_diagonal 2
15.228 ms (2 allocations: 7.62 MiB)

This is probably the lowest you can go in terms of allocations with this approach (mat in the function allocates 1000x999 Float64s which is basically the 7.6MB above)

a = [1 2 3
     4 5 6
     7 8 9]
b = Vector{Vector{eltype(a)}}(undef, size(a, 2))
for (i,c) in enumerate(eachcol(a))
    c2 = copy(c)
    deleteat!(c2, i)
    b[i] = c2
end
hcat(b...)

is faster, but the following will give you exactly what you want (but is slightly slower then the code above):

b = Vector{Vector{eltype(a)}}(undef, size(a, 2))
for (i,c) in enumerate(eachrow(a))
    c2 = copy(c)
    deleteat!(c2, i)
    b[i] = c2
end
hcat(b...)'

i.e.

2 3
4 6
7 8

but benchmark it and see what’s best (I have no idea).

@nilshg, @yakir12, thank you both, your methods both work and are much faster than mine!

Sorry, last one (I’m procrastinating…). This one is the prettiest I feel. No idea if it’s the fastest…

sz = LinearIndices(a)
n = size(a, 1)
k = [sz[i,i] for i in 1:n]
b = collect(vec(a'))
deleteat!(b, k)
reshape(b, n - 1, n)'

OK OK OK, I swear! Last one, best for last:

i = CartesianIndices(a)
k = filter(x -> x.I[1] ≠ x.I[2], i)
b = reshape(a'[k], n - 1, n)'

(I must come across as a total lunatic, I’m not, ask others here)

4 Likes

I can confirm that Yakir is not a total lunatic.

A quick benchmark:

# Original function
57.002 ms (19002 allocations: 75.87 MiB)
# nilshg function
  14.625 ms (2 allocations: 7.62 MiB)
# Yakir 1
  14.943 ms (2 allocations: 7.62 MiB)
# Yakir 2
  7.595 ms (8 allocations: 7.64 MiB)
# Yakir 3
  20.609 ms (13 allocations: 23.82 MiB)

Note that the second function that Yakir proposed returns an adjoint rather than an array - in all likelihood this is fine but depending on your use case you might have to collect it into a regular array, which can wipe out the gains:

# Yakir 2 plus collect
16.184 ms (10 allocations: 15.26 MiB)

Full benchmark code below for people to check whether I’ve screwed up the benchmarking somehow:


function remove_diagonal(x)
    mat = Array{Float64}(undef, size(x, 1), size(x, 1) - 1)
    for i = 1:size(x, 1)
        inds = setdiff(1:size(x, 1), i)
        mat[i, :] .= x[i, inds]
    end
    return mat
end


function remove_diagonal2(x)
    mat = Array{Float64}(undef, size(x, 1), size(x, 1) - 1)
    for i = 1:size(mat, 1)
        for j ∈ 1:size(mat, 2)
            if i > j
                mat[i, j] = x[i, j]
            else
                mat[i, j] = x[i, j+1]
            end
        end    
    end
    return mat
end


function remove_diagonal3(x)
    
    b = Vector{Vector{eltype(a)}}(undef, size(x, 2))
    for (i,c) in enumerate(eachrow(x))
        c2 = copy(c)
        deleteat!(c2, i)
        b[i] = c2
    end
    hcat(b...)'
end

function remove_diagonal4(x)
    sz = LinearIndices(x)
    n = size(x, 1)
    k = [sz[i,i] for i in 1:n]
    b = collect(vec(x'))
    deleteat!(b, k)
    reshape(b, n - 1, n)'
end

rd4_c(x) = collect(remove_diagonal4(x))

function remove_diagonal5(x)
    n = size(x, 1)
    i = CartesianIndices(x)
    k = filter(x -> x.I[1] ≠ x.I[2], i)
    b = reshape(x'[k], n - 1, n)'
end

using BenchmarkTools
xl = rand(1000, 1000)

println("Original function")
@btime remove_diagonal($xl)

println("nilshg function")
@btime remove_diagonal2($xl)

println("Yakir 1")
@btime remove_diagonal3($xl)

println("Yakir 2")
@btime remove_diagonal4($xl)

println("Yakir 3")
@btime remove_diagonal5($xl)

println("Yakir 2 plus collect")
@btime rd4_c($xl)
5 Likes