Enable Javascript in your browser and then refresh this page, for a much enhanced experience.
Fast; first rows with magic sums, then permutations solution in Clear category for Magic Square by kkkkk
from itertools import permutations
from itertools import combinations
def all_good_sums(data, magic_sum):
"""Return True if all columns and diagonals equal magic_sum, else False."""
left_to_right_diag = [[data[i][i] for i in range(len(data))]]
right_to_left_diag = [[row[-i-1] for i, row in enumerate(data)]]
for list_to_test in [zip(*data), left_to_right_diag, right_to_left_diag]:
for digits in list_to_test:
if sum(digits) != magic_sum:
return False
return True
def rows_with_magic_sum(data, missing_digits, magic_sum):
"""Return list of row combinations that all add up to magic sum."""
length = len(data)
# Based on the sum of the row alone (ignoring the order of numbers for
# now), supply various combination of numbers for the missing numbers in
# a row and record those combinations that add up to the magic sum.
combos_with_good_sums = []
stack = [(0, [], missing_digits)]
while stack:
row_idx, combo_list, remaining_nums = stack.pop()
# If this is the last row of numbers adding up to the magic
# numbers, save the rows for further processing by another function.
if row_idx == length:
combos_with_good_sums.append(combo_list)
continue
# For any numbers missing in the row, create combinations of unused
# numbers and check whether those numbers and the existing numbers
# add up to the magic sum. If so, save the numbers for this row on
# the stack and move onto the next row.
for combo in combinations(remaining_nums, data[row_idx].count(0)):
if sum(data[row_idx]) + sum(combo) == magic_sum:
stack.append((row_idx+1, combo_list + [list(combo)],
remaining_nums - set(combo)))
return combos_with_good_sums
def permutations_of_rows(data, row_combos, magic_sum):
"""Return valid magic square or None.
Use permutations of row combos to find a good magic square.
"""
length = len(data)
# We have combinations of numbers for a row that add up to a magic
# square, but we need to use permutations of those numbers to find
# columns and diagonals that add up to the magic sum.
stack = [(0, [])]
while stack:
row_idx, combo_list = stack.pop()
# If this is the last row of numbers, check if square has good
# sums. If so, we're done. Otherwise, continue to try the other
# permutations.
if row_idx == length:
if all_good_sums(combo_list, magic_sum):
return combo_list
continue
# Take the permutations for this row of numbers along with the
# numbers already present in the data and create a new square,
# row by row.
for combo in permutations(row_combos[row_idx]):
merged_row = data[row_idx][:]
combo_idx = 0
for idx, num in enumerate(merged_row):
if num == 0:
merged_row[idx] = combo[combo_idx]
combo_idx += 1
stack.append((row_idx + 1, combo_list + [merged_row]))
return None
def checkio(data):
"""Return completed magic square."""
# Calcuate the magic sum using the formula found on Wikipedia:
# (n * (n**2 + 1)) / 2
length = len(data)
magic_sum = (length * (length**2 + 1)) // 2
# Find list of all row combinations adding up to the magic sum.
digits = {x for row in data for x in row} - {0}
missing_digits = {x for x in range(1, length * length + 1)} - digits
good_row_combos = rows_with_magic_sum(data, missing_digits, magic_sum)
# For each row of potential substitutions, figure out which combination
# yields the magic sums column-wise and diagonally.
for row_combos in good_row_combos:
completed_square = permutations_of_rows(data, row_combos, magic_sum)
if completed_square:
return completed_square
return None
Oct. 8, 2020
