Mosaic diagonals

A mosaic wall is a rectangle of square tiles, each cut with a shade from 0 (palest) to 100 (darkest). The commission asks for a fade along every down-right diagonal, and the tiles are already cut.

The problem

The tiles are set in rows rows and cols columns. A diagonal is the run you get by starting on the top row or the left column and stepping one row down and one column right until you fall off the wall; every tile sits on exactly one.

The setter may swap tiles only within a diagonal — they are cut on a bias and will not seat anywhere else. Rearrange each diagonal so its shades run palest at the top-left end to darkest at the bottom-right, and return the wall.

Input. wall — a list of lists of integers, the shade of each tile, given row by row. Every row has the same length.

Output. The wall with each down-right diagonal sorted from pale to dark.

Example.

wall = [[5, 2, 9],          ->  [[4, 2, 9],
        [6, 4, 2],               [6, 4, 2],
        [3, 7, 4]]               [3, 7, 5]]

The main diagonal holds 5, 4, 4 and comes back as 4, 4, 5. The one above it is already 2, 2, the one below holds 6 then 7, and the corners are single tiles.

A second example, on a wall that is not square, where the diagonals have three different lengths:

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

The main diagonal 8, 9, 8 becomes 8, 8, 9, the pair 4, 3 becomes 3, 4, the pair 6, 5 becomes 5, 6, and the run of three 1s is unchanged.

Constraints.

• 1 <= rows, cols <= 100
• 0 <= wall[r][c] <= 100

Hints

Approach

Brute force

For each tile, walk its whole diagonal counting the tiles paler than it, and place it at that rank. Every tile walks up to min(rows, cols) neighbours — a million steps on a 100 by 100 wall, with fragile bookkeeping wherever two shades tie.

The insight

Every tile on a down-right diagonal has the same value of row − col, so that single number names the diagonal — bucket the tiles by it, sort each bucket, and lay them back down in the same traversal order.

Stepping down-right adds one to the row and one to the column, leaving row − col unchanged; every other step changes it, so the key partitions the wall exactly into diagonals. The write-back visits cells in the order the collection did, so each sorted bucket lands top-left to bottom-right.

Algorithm

1. For every cell, append its shade to the bucket keyed by row − col.
2. Sort each bucket in descending order.
3. Walk the cells in the same order again, popping the last value of each cell's bucket — the palest one left — into the cell.
4. Return the wall.

Complexity

Time O(rows · cols · log min(rows, cols)) — each tile is placed once, and a bucket of length L sorts in L log L. Space O(rows · cols) for the buckets.

Shades are capped at 100, so a counting sort per diagonal replaces the comparison sort: tally 101 slots and read them back in order. On a wall of short diagonals that is a real win; on long ones the two are close.

Solution

"""Mosaic diagonals — bucket tiles by row - col, sort each bucket, write back."""

from collections import defaultdict


def solve(wall):
    rows, cols = len(wall), len(wall[0])
    laid = [row[:] for row in wall]

    # invariant: row - col is constant along a down-right diagonal and differs
    # between diagonals, so these buckets partition the wall exactly.
    diagonals = defaultdict(list)
    for r in range(rows):
        for c in range(cols):
            diagonals[r - c].append(laid[r][c])

    # Sorted descending so that pop() hands back the palest tile left.
    for shades in diagonals.values():
        shades.sort(reverse=True)

    # The write-back walks the cells in the same order they were collected,
    # so each diagonal is filled top-left to bottom-right.
    for r in range(rows):
        for c in range(cols):
            laid[r][c] = diagonals[r - c].pop()
    return laid

TESTS = [
    (([[5, 2, 9], [6, 4, 2], [3, 7, 4]],), [[4, 2, 9], [6, 4, 2], [3, 7, 5]]),
    (([[8, 1, 4, 7], [6, 9, 1, 3], [2, 5, 8, 1]],), [[8, 1, 3, 7], [5, 8, 1, 4], [2, 6, 9, 1]]),
    (([[4, 3, 2, 1]],), [[4, 3, 2, 1]]),
    (([[4], [3], [2]],), [[4], [3], [2]]),
    (([[7, 7], [7, 7]],), [[7, 7], [7, 7]]),
    (([[9]],), [[9]]),
    (([[1, 2], [3, 4], [5, 6]],), [[1, 2], [3, 4], [5, 6]]),
]

Pitfalls

• Keying on row + col. That names the up-right anti-diagonal. The first example comes back as [[5, 2, 3], [6, 4, 2], [9, 7, 4]] — sorted along the wrong lines.
• Assuming the wall is square. Looping the column index over len(wall) drops the last column of the 3 by 4 example, and indexes out of range on a wall taller than it is wide.
• Starting every diagonal at wall[0][abs(d)]. Negative keys start on the top row, positive keys start on the left column, and treating them alike reads tiles off the wrong diagonal.

Variants

• Fairground scoreboard — there the comparison is the problem; here the key is.
• Order as a key — writing the order you want as a function of one element.