Nice way to iterate through a triangular matrix

I often need to loop through the upper/lower triangular section of a matrix and usually end up doing something like:

A = reshape(1:25,5,5)

for j in axes(A,2)
    for i in axes(A,1)
        if i>j
            @show A[i,j]
        end
    end
end

This is fine but can lead to code with way too many indents. You could get rid of the if statement, but its still very clunky and takes a second to figure out whats going on:

for j in axes(A,2)
    for i in j+1:size(A,1)
        @show A[i,j]
    end
end

I came up with simpler solution and wanted to share it. Basically Iterators does all the heavy lifting.

lowTriIter(ax1, ax2) = Iterators.filter(i -> first(i)>last(i),
                                        Iterators.product(ax1,ax2) )

lowTriIter(A::AbstractMatrix) = lowTriIter(axes(A)...)

Which leads to code that looks like this:

for (i,j) in lowTriIter(A)
    @show A[i,j]
end

It’s pretty versatile like you could change the comparison to include the diagonal and there’s no fussing around with adding one to indices. You could probably even get it to work with CartesianIndices if you wanted. Feel free to suggest improvements :grin:

2 Likes

Another way could be:

lowTriIter2(A::AbstractMatrix) = 
  Iterators.flatten(Iterators.map(j->(j,i), i+1:size(A,1)) for i=1:size(A,2))
julia> collect(lowTriIter2(A))
10-element Vector{Tuple{Int64, Int64}}:
 (2, 1)
 (3, 1)
 (4, 1)
 (5, 1)
 (3, 2)
 (4, 2)
 (5, 2)
 (4, 3)
 (5, 3)
 (5, 4)

This avoids looping over the skipped indices, and therefore might save a little time.

2 Likes

Oh I don’t think I would have ever thought of making a generator of maps, interesting. I guess you could skip right to the generator as well:

lowTriIter3(A::AbstractMatrix) =
    ((i,j) for j in axes(A,2) for i in j+1:size(A,1))

I do like that these methods skip over the additional iterations but I do wish they avoided the i+1:size(A,1) bit. Also I made my version more complicated than it need to be. I think I was worried about axes(A) having more than two iterators, but it won’t because A is already a matrix:

lowTriIter(A::AbstractMatrix) = 
    Iterators.filter(i -> first(i)>last(i), Iterators.product(axes(A)...))
1 Like

You know you can write both iterations on one line?

for j in axes(A,2), i in j+1:size(A,1)
    @show A[i,j]
end
8 Likes
3 Likes

I did, but that’s a fair point! It’s wild how many ways you can accomplish the same task.

I think I liked my version because it abstracted away the need to think about indices. Like if I wanted to change this nested loop to iterate over the upper triangle it would take me a minute to figure out versus just flipping a less-than sign. I get that’s just a preference though :slightly_smiling_face:

Ah i didn’t see this. Sorry for the repeated topics. That formula for the ut2ci(idx) function is both amazing and headache inducing :rofl:

Just some additional thoughts …

First, you could also use iterator comprehensions and the definition of your iterator looks just like your for loop:

lowertriag(A::AbstractMatrix) = ((i,j) for j in axes(A, 2) for i in axes(A, 1) if i > j)

for (i, j) in lowertriag(A)
    @show A[i, j]
end

Further, flatmap, i.e., flatten ∘ map, is a very powerful combinator and most other operations on iterables can be defined in terms of it:

mymap(f, xs) = Iterators.flatmap(x -> (f(x),), xs)
myfilter(pred, xs) = Iterators.flatmap(x -> if pred(x); (x,) else empty((x,)) end, xs)

In the functional programming jargon a la Haskell, flatmap is the monadic bind whereas map is the way more general functor instance for lists.

julia> (A[ci] for ci in CartesianIndices(A) if first(ci.I) ≤ last(ci.I)) |> collect
15-element Vector{Int64}:
  1
  6
  7
 11
 12
 13
 16
 17
 18
 19
 21
 22
 23
 24
 25

julia> (A[ci] for ci in CartesianIndices(A) if first(ci.I) ≥ last(ci.I)) |> collect
15-element Vector{Int64}:
  1
  2
  3
  4
  5
  7
  8
  9
 10
 13
 14
 15
 19
 20
 25
1 Like

Oh I did see that Iterators.flatmap was new to v1.9. It definitely looks useful but I’m not sure how I would apply it to this situation.

I’m not sure how this i > j comparison or i = j + 1 works on non-monotonic increasing array indices like StarWarsArrays?