[FIXED] Neater strategy to drop down the numbers in the 2D list

Issue

I have a problem. It is a 2D list of non-negative integers will be given like

0, 0, 2, 0, 1
0, 2, 1, 1, 0
3, 0, 2, 1, 0
0, 0, 0, 0, 0

I have to drop the numbers, number columns. e.g. drop down the 1's down 1 column, the 2's down 2 columns, the 3's down 3 columns, and so on. If the number can't be moved down enough, wrap it around the top. (e. g If there is a 3 in the second-to-last row, it should wrap around to the first row.) If two numbers map to the same slot, the biggest number takes that slot.

After this transformation the given matrix above will end up like:

0, 0, 2, 0, 0
3, 0, 0, 0, 1
0, 0, 2, 1, 0
0, 2, 0, 1, 0

Here's my trivial solution to the problem (Assumes a list l is pre-set):

new = [[0] * len(l[0]) for _ in range(len(l))]
idx = sorted([((n + x) % len(l), m, x) for n, y in enumerate(l) for m, x in enumerate(y)], key=lambda e: e[2])
for x, y, z in idx:
  new[x][y] = z
print(new)

The strategy is:

• Build a list new with 0s of the shape of l
• Save the new indices of each number in l and each number as tuple pairs in idx
• Sort idx by each number
• Assign indices from idx to the respective numbers to new list
• Print new

I am not satisfied with this strategy. Is there a neater/better way to do this? I can use numpy.

Solution

Let's say you have

a = np.array([
    [0,0,2,0,1],
    [0,2,1,1,0],
    [3,0,2,1,0],
    [0,0,0,0,0]])

You can get the locations of the elements with np.where or np.nonzero:

r, c = np.nonzero(a)

And the elements themselves with the index:

v = a[r, c]

Incrementing the row is simple now:

new_r = (r + v) % a.shape[0]

To settle collisions, sort the arrays so that large values come last:

i = v.argsort()

Now you can assign to a fresh matrix of zeros directly:

result = np.zeros_like(a)
result[new_r[i], c[i]] = v[i]

The result is

[[0 0 2 0 0]
 [3 0 0 0 1]
 [0 0 2 1 0]
 [0 2 0 1 0]]

Answered By - Mad Physicist