Get all the diagonals in a matrix/list of lists in Python
I came across another interesting solution to this issue. The row, column, forward, and backward diagonal can all be immediately discovered by looking at a combination of x and y.
Column = x Row = y F-Diag = x+y B-Diag = x-y B-Diag` = x-y-MIN
| 0 1 2 | 0 1 2 | 0 1 2 | 0 1 2 | 0 1 2
--|--------- --|--------- --|--------- --|--------- --|---------
0 | 0 1 2 0 | 0 0 0 0 | 0 1 2 0 | 0 1 2 0 | 2 3 4
1 | 0 1 2 1 | 1 1 1 1 | 1 2 3 1 |-1 0 1 1 | 1 2 3
2 | 0 1 2 2 | 2 2 2 2 | 2 3 4 2 |-2 -1 0 2 | 0 1 2
From the diagram you can see that each diagonal and axis is uniquely identifiable using these equations. Take each unique number from each table and create a container for that identifier.
Note that the backward diagonals have been offset to start at a zero index, and that the length of forward diagonals is always equal to the length of backward diagonals.
test = [[1,2,3],[4,5,6],[7,8,9],[10,11,12]]
max_col = len(test[0])
max_row = len(test)
cols = [[] for _ in range(max_col)]
rows = [[] for _ in range(max_row)]
fdiag = [[] for _ in range(max_row + max_col - 1)]
bdiag = [[] for _ in range(len(fdiag))]
min_bdiag = -max_row + 1
for x in range(max_col):
for y in range(max_row):
cols[x].append(test[y][x])
rows[y].append(test[y][x])
fdiag[x+y].append(test[y][x])
bdiag[x-y-min_bdiag].append(test[y][x])
print(cols)
print(rows)
print(fdiag)
print(bdiag)
Which will print
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
[[1, 4, 7, 10], [2, 5, 8, 11], [3, 6, 9, 12]]
[[1], [2, 4], [3, 5, 7], [6, 8, 10], [9, 11], [12]]
[[10], [7, 11], [4, 8, 12], [1, 5, 9], [2, 6], [3]]
Using a defaultdict and a lambda, this can be generalized further:
from collections import defaultdict
def groups(data, func):
grouping = defaultdict(list)
for y in range(len(test)):
for x in range(len(test[y])):
grouping[func(x, y)].append(data[y][x])
return list(map(grouping.get, sorted(grouping)))
test = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
cols = groups(test, lambda x, y: x)
rows = groups(test, lambda x, y: y)
fdiag = groups(test, lambda x, y: x + y)
bdiag = groups(test, lambda x, y: x - y)
There are probably better ways to do it in numpy than below, but I'm not too familiar with it yet:
import numpy as np
matrix = np.array(
[[-2, 5, 3, 2],
[ 9, -6, 5, 1],
[ 3, 2, 7, 3],
[-1, 8, -4, 8]])
diags = [matrix[::-1,:].diagonal(i) for i in range(-3,4)]
diags.extend(matrix.diagonal(i) for i in range(3,-4,-1))
print [n.tolist() for n in diags]
Output
[[-2], [9, 5], [3, -6, 3], [-1, 2, 5, 2], [8, 7, 1], [-4, 3], [8], [2], [3, 1], [5, 5, 3], [-2, -6, 7, 8], [9, 2, -4], [3, 8], [-1]]
Edit: Updated to generalize for any matrix size.
import numpy as np
# Alter dimensions as needed
x,y = 3,4
# create a default array of specified dimensions
a = np.arange(x*y).reshape(x,y)
print a
print
# a.diagonal returns the top-left-to-lower-right diagonal "i"
# according to this diagram:
#
# 0 1 2 3 4 ...
# -1 0 1 2 3
# -2 -1 0 1 2
# -3 -2 -1 0 1
# :
#
# You wanted lower-left-to-upper-right and upper-left-to-lower-right diagonals.
#
# The syntax a[slice,slice] returns a new array with elements from the sliced ranges,
# where "slice" is Python's [start[:stop[:step]] format.
# "::-1" returns the rows in reverse. ":" returns the columns as is,
# effectively vertically mirroring the original array so the wanted diagonals are
# lower-right-to-uppper-left.
#
# Then a list comprehension is used to collect all the diagonals. The range
# is -x+1 to y (exclusive of y), so for a matrix like the example above
# (x,y) = (4,5) = -3 to 4.
diags = [a[::-1,:].diagonal(i) for i in range(-a.shape[0]+1,a.shape[1])]
# Now back to the original array to get the upper-left-to-lower-right diagonals,
# starting from the right, so the range needed for shape (x,y) was y-1 to -x+1 descending.
diags.extend(a.diagonal(i) for i in range(a.shape[1]-1,-a.shape[0],-1))
# Another list comp to convert back to Python lists from numpy arrays,
# so it prints what you requested.
print [n.tolist() for n in diags]
Output
[[ 0 1 2 3]
[ 4 5 6 7]
[ 8 9 10 11]]
[[0], [4, 1], [8, 5, 2], [9, 6, 3], [10, 7], [11], [3], [2, 7], [1, 6, 11], [0, 5, 10], [4, 9], [8]]