DEV Community

Discussion on: Daily Challenge #105 - High-Sum Matrix Drop

Collapse
 
sardok profile image
sinx

My solution using queue. Should be o(n).


import string

def get_steps(i, j, row, col):
    left = j - 1
    right = j + 1
    down = i + 1

    if down >= row:
        return []

    steps = []
    if left >= 0:
        steps.append((down, left))

    steps.append((down, j))

    if right < col:
        steps.append((down, right))

    return steps

def make_and_fill(row, col, fill):
    ret = []
    for _ in range(row):
        row_ = [fill] * col
        ret.append(row_)

    return ret

def find_max_index(xs):
    max_ = max(xs)
    return xs.index(max_)

def get_path(i, j, backtraces):
    ret = []
    queue = [(i, j)]

    while queue:
        pos = queue.pop(0)
        if pos:
            ret.append(pos)
            i, j = pos
            backtrace = backtraces[i][j]
            queue.append(backtrace)

    ret.reverse()
    ret = [f'{string.ascii_letters[j]}/{i}' for (i, j) in ret]
    return ', '.join(ret)


def solve(xs):
    row, col = len(xs), len(xs[0])
    sums = make_and_fill(row, col, 0)
    sums[0][:] = xs[0]
    backtraces = make_and_fill(row, col, 0)
    queue = [(0, col//2)]

    while queue:
        pos = queue.pop(0)
        i, j = pos
        for step in get_steps(i, j, row, col):
            step_i, step_j = step
            acc = sums[i][j] + xs[step_i][step_j]
            if sums[step_i][step_j] < acc:
                queue.append((step_i, step_j))
                sums[step_i][step_j] = acc
                backtraces[step_i][step_j] = pos

    max_index = find_max_index(sums[-1])
    print(sums[-1][max_index])

    path = get_path(row-1, max_index, backtraces)
    print(path)


if __name__ == '__main__':
    matrix = [
        [2, 1, 4, 5, 3],
        [5, 2, 1, 4, 3],
        [1, 3, 2, 5, 4],
        [1, 5, 2, 3, 4],
        [3, 2, 1, 5, 4],
    ]
    solve(matrix)