sandpiles.py

class Sandpile:
    def __init__(self, size, data=None):
        self.size = size
        
        if data:
            assert(isinstance(data, list))
            assert(len(data) == size)
            
            for row in data:
                assert(isinstance(row, list))
                assert(len(row) == size)
                
                for value in row:
                    assert(isinstance(value, int))
            
            self.data = data
        else:
            self.data = [[0] * size for _ in range(0, size)]
            
    def __add__(self, other):
        def requires_toppling(data):
            result = False
            for row in data:
                for value in row:
                    if value > 3:
                        result = True
            return result
        
        def get_neighbours(row_index, column_index, size):
            results = []
            for row_offset, column_offset in [(0, -1), (1, 0), (0, 1), (-1, 0)]:
                new_row_index = row_index + row_offset
                new_column_index = column_index + column_offset
                
                is_row_valid = new_row_index >= 0 and new_row_index < size
                is_column_valid = new_column_index >= 0 and new_column_index < size
                
                if is_row_valid and is_column_valid:
                    results.append((new_row_index, new_column_index))
            return results
        
        assert(self.size == other.size)
        
        size = self.size        
        
        # clone data
        new_data = [row[:] for row in self.data]
    
        for row_index, row in enumerate(other.data):
            for column_index, value in enumerate(row):
                new_data[row_index][column_index] += value
        
        while requires_toppling(new_data):            
            toppled_data = [row[:] for row in new_data]
            
            for row_index, row in enumerate(toppled_data):
                for column_index, value in enumerate(row):
                    if value > 3:
                        toppled_data[row_index][column_index] = value % 4
                        
                        amount_to_topple = value // 4
                        
                        neighbours = get_neighbours(row_index, column_index, size)
                        for neighbour_row_index, neighbour_column_index in neighbours:
                            toppled_data[neighbour_row_index][neighbour_column_index] += amount_to_topple
            
            new_data = toppled_data
        
        return Sandpile(size, data=new_data)
    
    def _data_to_string(self, data):
        maximum_value_length = 0
        for row in data:
            for value in row:
                maximum_value_length = max(maximum_value_length, len(str(value)))
        
        row_strs = []
        for row in data:
            value_strs = []
            for value in row:
                # pad the value out to the required length
                value_str = '{0: >{1}}'.format(value, maximum_value_length)
                value_strs.append(value_str)
            row_strs.append(' '.join(value_strs))
        
        return '\n'.join(row_strs)
            
    def __str__(self):
        return self._data_to_string(self.data)

a_data = [
    [2, 1, 2],
    [1, 0, 1],
    [2, 1, 2]
]
a = Sandpile(3, data=a_data)

b_data = [
    [2, 2, 2],
    [2, 2, 2],
    [2, 2, 2]
]
b = Sandpile(3, data=b_data)

print(a + b)