[FIXED] Converting a list into a matrix by diagonally filling up the elements

Issue

I have a sorted list which looks as below:

mylist = [-2, -2, 1, 1, 4, 4, 3, 3, 3]

The list is sorted in ascending order based upon the number of times it appears. In case of a tie, the list is sorted based upon the values.

I need to convert this list into a square matrix of equal chunks(3*3 in this case) such that the numbers are placed "diagonally" starting from the bottom right corner.

The general case is to divide the list in equal chunks.

Desired Output:

res = [[3, 3, 4],
      [3, 4, 1],
      [1, -2, -2]]

I have written the below code but still not able to get the desired output:

def create_matrix(lst, n):
    for i in range(0, len(lst), n):
        print(i)
        yield lst[i: i+n]

m = create_matrix(mylist, 3)

print(list(m))

One solution could be to place pairs in queue/stack and then pop as needed.

Solution

I assume you want to iterate the output matrix this way (NxN = 5x5 example, 0->24 order):

[[ 0,  1,  3,  6, 10],
 [ 2,  4,  7, 11, 15],
 [ 5,  8, 12, 16, 19],
 [ 9, 13, 17, 20, 22],
 [14, 18, 21, 23, 24]]

For each cell, the coordinates (i,j) have their sum equal to the number of the diagonal (k) from top-left to bottom right (2*N-1 diagonals in total)

For the N first diagonals, the first item has i=0, the following ones i=k-N where k is the diagonal number.

The last item has i=k with a maximum of N.

j = k-i

This gives us the following algorithm to iterate the cells in order:

import math

mylist = [-2, -2, 1, 1, 4, 4, 3, 3, 3]

N = int(math.sqrt(len(mylist))) # 3

out = [[None for _ in range(N)]
       for _ in range(N)]

r = reversed(mylist)

for k in range(2*N-1):
    start = 0 if k<N else k-N+1
    stop = min(k, N-1)
    for i in range(start, stop+1):
        out[i][k-i] = next(r)

print(out)

Output:

[[3, 3, 4],
 [3, 4, 1],
 [1, -2, -2]]

alternative approach

If you want to tackle the problem the other way around and generate the items in order of the rows, then columns, you can use a generator:

def yield_diag(lst, N):
    for I in range(N):
        # sum of successive diagonal lengths
        i = I*(I+1)//2
        j = 0
        for J in range(I, N+I):
            # step between columns
            j += min(J, N)-max(0, J-N+1)
            yield lst[i+j]

# 5x5 example
list(yield_diag(list(range(25)), 5))
# [0, 1, 3, 6, 10, 2, 4, 7, 11, 15, 5, 8, 12, 16, 19, 9, 13, 17, 20, 22, 14, 18, 21, 23, 24]

OP example:

list(yield_diag(mylist[::-1], 3))

# [3, 3, 4, 3, 4, 1, 1, -2, -2]

As 2D:

N = 3

it = yield_diag(mylist[::-1], N)

[[next(it) for _ in range(N)]
 for _ in range(N)]

Output:

[[3, 3, 4],
 [3, 4, 1],
 [1, -2, -2]]

Answered By - mozway