Get all the diagonals in a matrix/list of lists in Python

I'm looking for a Pythonic way to get all the diagonals of a (square) matrix, represented as a list of lists.

Suppose I have the following matrix:

matrix = [[-2,  5,开发者_如何学JAVA  3,  2],
          [ 9, -6,  5,  1],
          [ 3,  2,  7,  3],
          [-1,  8, -4,  8]]

Then the large diagonals are easy:

l = len(matrix[0])
print([matrix[i][i] for i in range(l)])              # [-2, -6, 7,  8]
print([matrix[l-1-i][i] for i in range(l-1,-1,-1)])  # [ 2,  5, 2, -1]

But I have trouble coming up with a way to generate all the diagonals. The output I'm looking for is:

[[-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]]

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]]

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)

Start with the diagonals that slope up-and-right.

If (x,y) is a rectangular coordinate inside the matrix, you want to transform to/from a coordinate scheme (p,q), where p is the number of the diagonal and q is the index along the diagonal. (So p=0 is the [-2] diagonal, p=1 is the [9,5] diagonal, p=2 is the [3,-6,3] diagonal, and so on.)

To transform a (p,q) into an (x,y), you can use:

x = q
y = p - q

Try plugging in values of p and q to see how this is working.

Now you just loop... For p from 0 to 2N-1, and q from max(0, p-N+1) to min(p, N-1). Transform p,q to x,y and print.

Then for the other diagonals, repeat the loops but use a different transformation:

x = N - 1 - q
y = p - q

(This effectively just flips the matrix left-right.)

Sorry I did not actually code this in Python. :-)

This is for Moe, who asked a similar question.

I start off by making simple functions to copy rows or columns of any rectangular matrix.

def get_rows(grid):
    return [[c for c in r] for r in grid]

def get_cols(grid):
    return zip(*grid)

With these two functions I then get the diagonals by adding an increasing/decreasing buffer to the start/end of each row. I then get the columns of this buffered grid, then remove the buffer on each column afterwards. ie)

1 2 3    |X|X|1|2|3|    | | |1|2|3|
4 5 6 => |X|4|5|6|X| => | |4|5|6| | => [[7],[4,8],[1,5,9],[2,6],[3]]
7 8 9    |7|8|9|X|X|    |7|8|9| | |

.

def get_backward_diagonals(grid):
    b = [None] * (len(grid) - 1)
    grid = [b[i:] + r + b[:i] for i, r in enumerate(get_rows(grid))]
    return [[c for c in r if c is not None] for r in get_cols(grid)]

def get_forward_diagonals(grid):
    b = [None] * (len(grid) - 1)
    grid = [b[:i] + r + b[i:] for i, r in enumerate(get_rows(grid))]
    return [[c for c in r if c is not None] for r in get_cols(grid)]

I ended up reinventing this wheel recently. Here's an easy-to-reuse/extend method to find the diagonals in a square list-of-lists:

def get_diagonals(grid, bltr = True):
  dim = len(grid)
  assert dim == len(grid[0])
  return_grid = [[] for total in xrange(2 * len(grid) - 1)]
  for row in xrange(len(grid)):
    for col in xrange(len(grid[row])):
      if bltr: return_grid[row + col].append(grid[col][row])
      else:    return_grid[col - row + (dim - 1)].append(grid[row][col])
  return return_grid

Assuming list indices:

00 01 02 03

10 11 12 13

20 21 22 23

30 31 32 33

then setting bltr = True (the default), returns the diagonals from bottom-left to top-right, i.e.

00           # row + col == 0
10 01        # row + col == 1
20 11 02     # row + col == 2
30 21 12 03  # row + col == 3
31 22 13     # row + col == 4
32 23        # row + col == 5
33           # row + col == 6

setting bltr = False, returns the diagonals from bottom-left to top-right, i.e.

30            # (col - row) == -3
20 31         # (col - row) == -2
10 21 32      # (col - row) == -1
00 11 22 33   # (col - row) == 0
01 12 23      # (col - row) == +1
02 13         # (col - row) == +2
03            # (col - row) == +3

Here's a runnable version using OP's input matrix.

I guess there's an easier way to do this now. (But only use this if you are already familiar with the above answers).

from collections import defaultdict

There's this method called defaultdict which is imported from the collections module, is used to create dictionaries if you don't know the key you are going to have.

We use this in these situations:

• If you don't know the key but want to assign some value to a particular key.
• Normal dictionary raises keyerror if the key is not present in the dictionary. But this won't ( you can assign some function to it if you want)

After Importing, you can run the following code and check.

rows,cols = 3,3
matrix = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]

diagonal1 = defaultdict(list) # For the top right to bottom left
diagonal2 = defaultdict(list) # For the top left to bottom right
for i in range(rows):
    for j in range(cols):
        diagonal1[i-j].append(matrix[i][j])
        diagonal2[i+j].append(matrix[i][j])
print(diagonal1,'\n',diagonal2)

The list parameter will create a list of values for that particular key.

The output is as follows:

defaultdict(<class 'list'>, {0: [1, 5, 9], -1: [2, 6], -2: [3], 1: [4, 8], 2: [7]}) 
defaultdict(<class 'list'>, {0: [1], 1: [2, 4], 2: [3, 5, 7], 3: [6, 8], 4: [9]})

Now you can use both the diagonals as you want.

To know more about defaultdict use this link : Click here

This only works for matricies of equal width and height. But it also doesn't rely on any third parties.

matrix = [[11, 2, 4],[4, 5, 6],[10, 8, -12]]

# only works for diagnoals of equal width and height
def forward_diagonal(matrix):
    if not isinstance(matrix, list):
        raise TypeError("Must be of type list")

    results = []
    x = 0
    for k, row in enumerate(matrix):
        # next diag is (x + 1, y + 1)
        for i, elm in enumerate(row):

            if i == 0 and k == 0:
                results.append(elm)
                break
            if (x + 1 == i):
                results.append(elm)
                x = i
                break
    return results

print 'forward diagnoals', forward_diagonal(matrix)

Code based on Nemo's answer above:

def print_diagonals(matrix):
    n = len(matrix)
    diagonals_1 = []  # lower-left-to-upper-right diagonals
    diagonals_2 = []  # upper-left-to-lower-right diagonals
    for p in range(2*n-1):
        diagonals_1.append([matrix[p-q][q] for q in range(max(0, p - n + 1), min(p, n - 1) + 1)])
        diagonals_2.append([matrix[n-p+q-1][q] for q in range(max(0, p - n + 1), min(p, n - 1) + 1)])
    print("lower-left-to-upper-right diagonals: ", diagonals_1)
    print("upper-left-to-lower-right diagonals: ", diagonals_2)


print_diagonals([
    [1, 2, 1, 1],
    [1, 1, 4, 1],
    [1, 3, 1, 6],
    [1, 7, 2, 5],
])

lower-left-to-upper-right diagonals:  [[1], [1, 2], [1, 1, 1], [1, 3, 4, 1], [7, 1, 1], [2, 6], [5]]
upper-left-to-lower-right diagonals:  [[1], [1, 7], [1, 3, 2], [1, 1, 1, 5], [2, 4, 6], [1, 1], [1]]

Pythonic approach

For a pure Python implementation I would suggest to work in 1D.

W, H = len(mat[0]), len(mat) 
idx = range(W-1) + range(W-1, W*H, W)
rng = range(1, W) + range(H, 0, -1)
rng = map(lambda x: x if (x < min(W, H)) else min(W, H), rng)
dia = [[i + (W-1) * m for m in xrange(r)] for i, r in zip(idx, rng)]

Here dia returns a list of indices for each diagonal. To retrieve the corresponding values:

arr = [e for row in mat for e in row] #Flatten the matrix
for d in dia:
    print [arr[e] for e in d][::-1]

[-2]
[9, 5]
[3, -6, 3]
[-1, 2, 5, 2]
[8, 7, 1]
[-4, 3]
[8]

If you want to return the values in the opposite direction:

arr2 = [e for row in zip(*mat[::-1]) for e in row] #Flatten and rotate the matrix by 90°
for d in dia[::-1]:
    print [arr2[e] for e in d]

[2]
[3, 1]
[5, 5, 3]
[-2, -6, 7, 8]
[9, 2, -4]
[3, 8]
[-1]

Numpy approach

tril = [np.flip(np.fliplr(mat).diagonal(n)) for n in xrange(mat.shape[0])][::-1]
trir = [np.flipud(mat).diagonal(n) for n in xrange(1, mat.shape[0])]
dia = tril + trir

[array([-2]),
 array([9, 5]),
 array([ 3, -6,  3]),
 array([-1,  2,  5,  2]),
 array([8, 7, 1]),
 array([-4,  3]),
 array([8])]

Try this :

import numpy as np
matrix = [[-2,  5,  3,  2],
          [ 9, -6,  5,  1],
          [ 3,  2,  7,  3],
          [-1,  8, -4,  8]]

matrix = np.array(matrix)
matrix = np.flipud(matrix)
a = matrix.shape[0]
list_ = [np.diag(matrix, k=i).tolist() for i in range(-a+1,a)]
print(list_)

Output :

[[-2], [9, 5], [3, -6, 3], [-1, 2, 5, 2], [8, 7, 1], [-4, 3], [8]]

Try using dict

mat = [[-2,  5,  3,  2],
      [ 9, -6,  5,  1],
      [ 3,  2,  7,  3],
      [-1,  8, -4,  8]]
dct = dict()
for i in range(len(mat)-1,-len(mat[0]),-1):
    dct[i] = []
for i in range(len(mat)):
    for j in range(len(mat[0])):
       dct[i-j].append(mat[i][j])
print(dct)

Output:

{3: [-1], 2: [3, 8], 1: [9, 2, -4], 0: [-2, -6, 7, 8], -1: [5, 5, 3], -2: [3, 1], -3: [2]}

Using itertools

matrix = [[-2,  5,  3,  2],
      [ 9, -6,  5,  1],
      [ 3,  2,  7,  3],
      [-1,  8, -4,  8]]

import itertools as it

def show_diagonals(alist):

    # get row/col lenght
    a = len(alist)

    # creating a fliped matrix
    rlist = []
    for r in alist:
        new = r.copy()
        new.reverse()
        rlist.append(new)

    flatten_list = list(it.chain.from_iterable(alist)) 
    flatten_rlist = list(it.chain.from_iterable(rlist)) 
    b = len(flatten_list)
    first_diag = list(it.islice(flatten_list, 0, b+1, a+1))
    second_diag = list(it.islice(flatten_rlist, 0, b+1, a+1))
    return first_diag, second_diag

a, b = show_diagonals(matrix)

Using some numpy-fu to get the main diagonal:

import numpy as np  
r = np.arange(36) 
r.resize((6, 6)) 
print(r) 
r = r.reshape(len(r)**2)[::len(r)+1] 
print(r)

Prints:

[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]
 [18 19 20 21 22 23]
 [24 25 26 27 28 29]
 [30 31 32 33 34 35]]

[ 0  7 14 21 28 35]

From here : np.Diagonal

 np.diagonal(matrix)