From a bit of reasoning we have that for each diagonal in a matrix the difference between the column number (or x
in what follows) and the row number (y
) is constant. We can use as our data structure a collections.defaultdict
that is autoinitialised as an empty list and have a loop on all the matrix elements, finding the diagonal to which each element belongs and listing those elements in our defaultdict
.
def getdiags(matrix, nr, nc):
from collections import defaultdict
d = defaultdict(list)
for y in range(nr):
for x in range(nc):
d[x-y].append(matrix[y][x])
return d
We can profit also from an utility function to present in order our results
def printdiags(diags):
for i in sorted(diags):
print diags[i]
Finally we test our things in IPython
In [1]: def getdiags(matrix, nr, nc):
from collections import defaultdict
d = defaultdict(list)
for y in range(nr):
for x in range(nc):
d[x-y].append(matrix[y][x])
return d
...:
In [2]: def printdiags(diags):
for i in sorted(diags):
print diags[i]
...:
In [3]: from pprint import pprint
In [4]: m = [["%3.3d"%(r*100+c) for c in range(5)] for r in range(4)]
In [5]: pprint(m)
[['000', '001', '002', '003', '004'],
['100', '101', '102', '103', '104'],
['200', '201', '202', '203', '204'],
['300', '301', '302', '303', '304']]
In [6]: diags = getdiags(m, 4, 5)
In [7]: printdiags(diags)
['300']
['200', '301']
['100', '201', '302']
['000', '101', '202', '303']
['001', '102', '203', '304']
['002', '103', '204']
['003', '104']
['004']
In [8]:
That's all
Addendum
In a late edit, the OP stated that their input is a list of strings of equal length and that the answer is seeked in terms of a list of diagonals.
My getdiags
above works as well with the new alternative data structure and obtaining the seeked list is very simple:
def listofdiags(diags):
return [''.join(diags[i]) for i in sorted(diags)]
Of course this conversion can be implemented also inside getdiagonals
but that is left as an exercise.
Look, no defaultdict
# List of Strings from Diagonals
def lsd(m, nr, nc):
res = ["" for i in range(nr+nc-1)]
for r in range(nr):
for c in range(nc):
i = c-r+nr-1
res[i] = res[i]+m[r][c]
return res
pprint(lsd(m, 4, 5))
# ['300',
# '200301',
# '100201302',
# '000101202303',
# '001102203304',
# '002103204',
# '003104',
# '004']
yield
ing the solution
The following is less efficient but for the sake of completeness, here it is:
def enumdiags(m, nr, nc):
for i in range(nr+nc-1):
s = ""
for r in range(nr):
for c in range(nc):
if c-r+nr-1 == i : s = s+m[r][c]
yield i, s
for i, s in enumdiags(m, 4, 5):
print i, s
# 0 300
# 1 200301
# 2 100201302
# 3 000101202303
# 4 001102203304
# 5 002103204
# 6 003104
# 7 004