In Python, is there a way to “move/shift down” elements from one array to another based on another array?

Say I have nested list containing data (note: each nested list is always the same lenght), like this:

[
    [0, 1, 2], 
    [3, 4, 5],
    [6, 7, 8]
]

and another list (with the same lenght of each nested list in the above mentioned data list), like this:

[0, 1, 2]

Is it possible to “shift/move down” the elements in the nested lists by the amount provided in the second list, create additional lists at the end based on the max shift/move amount (in this case 2), and fill the “empty” spaces with 0, so the outcome would be this:

[
    [0, 0, 0],
    [3, 1, 0],
    [6, 4, 2],
    [0, 7, 4],
    [0, 0, 8]
]

So in this case: the first elements in each list don’t get moved down, the second elements get moved down 1 space, and the third elements get moved down 2 spaces. All “empty” spaces just get filled in with 0.

Is this somewhat decently achievable in bare Python (3.9)? I’ve also looked at the NumPy documentation, but haven’t really found anything resembling this particular problem, so if possible with NumPy (or any other package), that’s good too. Thanks.

Answer

Here’s a solution that doesn’t use any packages:

def transpose(m):
    return [list(x) for x in zip(*m)]

def apply_shifts(m, shifts):
    m_t = transpose(m)
    for i in range(len(shifts)):
        padding = (max(shifts) - shifts[i])
        m_t[i] = shifts[i] * [0] + m_t[i] + padding * [0]
    return transpose(m_t)

If we try it using your example, we can see the function behaves correctly:

>>> m = [
...     [0, 1, 2],
...     [3, 4, 5],
...     [6, 7, 8],
... ]
>>> shifts = [0, 1, 2]
>>> apply_shifts(m, shifts)
[
    [0, 0, 0], 
    [3, 1, 0], 
    [6, 4, 2], 
    [0, 7, 5], 
    [0, 0, 8]
]