Table scanner - answer to interview question

table.py

class Table:
    """A table of numbers"""

    def __init__(self, *args, **kwargs):
        if len(args) == 3:
            self.from_vals(*args)
        elif len(args) == 2:
            x, y = args
            if kwargs['mode'] == 'random':
                vals = [random.randint(0, 9) for _ in range(x * y)]
                self.from_vals(x, y, vals)
            elif kwargs['mode'] == 'blank':
                vals = [0 for _ in range(x * y)]
                self.from_vals(x, y, vals)

    def from_vals(self, x, y, vals):
        """Construct table from 1-dimensional array"""
        self.vals = [[vals[_x + _y * x] for _y in range(y)] for _x in range(x)]
        self.x = x
        self.y = y

    def __str__(self):
        """Print as formatted table"""
        def _str(item):
            if item is None:
                return '■'
            else:
                return str(item)

        max_width = max(
            len(_str(self[x, y]))
            for y in range(self.y)
            for x in range(self.x)
        )

        def row_to_str_iter(y):
            for x in range(self.x):
                str_item = _str(self[x, y])
                padding = " " * (max_width - len(str_item))
                yield padding + str_item

        res = ""
        for y in range(self.y):
            res += ' | '.join(row_to_str_iter(y))
            res += "\n"

        return res

    def __getitem__(self, item):
        """Allows for read access using `self[x, y]`"""
        x, y = item
        return self.vals[x][y]

    def __setitem__(self, key, value):
        """Allows for write access using `self[x, y] = value`"""
        x, y = key
        self.vals[x][y] = value

    def rows(self):
        """Iterate through rows"""
        for y in range(self.y):
            yield [self[x, y] for x in range(self.x)]

    def columns(self):
        """Iterate through columns"""
        for x in range(self.x):
            yield [self[x, y] for y in range(self.y)]

    def copy(self):
        """Deep copy"""
        vals = [self[x, y] for y in range(self.y) for x in range(self.x)]
        return Table(self.x, self.y, vals)

table_scanner.py

import itertools
import random
from table import Table

def right_sum_rows(table):
    """Sum rows from left to right. Each result cell is the sum of all input cells to the right."""
    sum_table = []
    for row in table.rows():
        sum = 0
        for x in range(table.x):
            sum_table.append(sum)
            sum += row[x]
    return Table(table.x, table.y, sum_table)


def find_row_candidates(table, right_sum):
    """Sum rows right to left, to find all row candidates. Returns a bitset of candidates."""
    candidates = 0
    for y, (row, row_sum) in enumerate(zip(table.rows(), right_sum.rows())):
        sum = 0
        for x in reversed(range(table.x)):
            if sum == row_sum[x]:
                candidates |= 1 << (x + y * table.x)
            sum += row[x]

    return candidates


def down_sum_columns(table):
    """Sum columns from top to bottom. Each result cell is the sum of all input cells above it."""
    sum_table = Table(table.x, table.y, mode='blank')
    for x, column in enumerate(table.columns()):
        sum = 0
        for y, item in enumerate(column):
            sum_table[x, y] = sum
            sum += item
    return sum_table


def find_column_candidates(table, down_sum):
    """Sum rows bottom to top, to find all column candidates. Returns a bitset of candidates."""
    candidates = 0
    for x, (column, column_sum) in enumerate(zip(table.columns(), down_sum.columns())):
        sum = 0
        for y in reversed(range(table.y)):
            if sum == column_sum[y]:
                candidates |= 1 << (x + y * table.x)
            sum += column[y]

    return candidates


# table = Table(5, 5, mode='blank')
table = Table(6, 5,
              [0, 6, 1, 0, 0, 0,
               0, 0, 5, 0, 0, 0,
               0, 1, 0, 0, 0, 0,
               9, 5, 0, 0, 8, 6,
               0, 1, 0, 0, 0, 0]
              )

print('Table')
print(table)
right_sum = right_sum_rows(table)

print('Left to right sum')
print(right_sum)
candidates = find_row_candidates(table, right_sum)


def print_candidates(candidates):
    candidates_table = table.copy()
    i = 0
    for y in range(table.y):
        for x in range(table.x):
            if (candidates >> i) & 1 == 1:
                candidates_table[x, y] = None
            i += 1

    print('Candidates')
    print(candidates_table)

print('Row candidates')
print_candidates(candidates)

down_sum = down_sum_columns(table)
print('Top to bottom sum')
print(down_sum)

print('Column candidates')
print_candidates(find_column_candidates(table, down_sum))

candidates &= find_column_candidates(table, down_sum)
print('Final answer')
print_candidates(candidates)