-
Notifications
You must be signed in to change notification settings - Fork 17
/
magic-square-forming.py
92 lines (68 loc) · 2.32 KB
/
magic-square-forming.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#!/bin/python3
import math
import os
import random
import re
import sys
from typing import List
#
# Complete the 'formingMagicSquare' function below.
#
# The function is expected to return an INTEGER.
# The function accepts 2D_INTEGER_ARRAY s as parameter.
#
def formingMagicSquare(square: List[List[int]]) -> int:
rows, cols = len(square), len(square[0]) if square else 0
magick_square = [[0] * cols for _ in range(rows)]
max_value = rows * cols
def validate() -> bool:
cols_sums: List[int] = [0] * cols
rows_sums: List[int] = [0] * rows
diag1, diag2 = 0, 0
for row in range(rows):
for col in range(cols):
cols_sums[col] += magick_square[row][col]
rows_sums[row] += magick_square[row][col]
if row == col:
diag1 += magick_square[row][col]
if (cols - col - 1) == row:
diag2 += magick_square[row][col]
# It is a square
return (
cols_sums == rows_sums
and len(set(cols_sums)) == 1
and cols_sums[0] == diag1 == diag2
)
def diff(left: List[List[int]], right: List[List[int]]) -> int:
_diff = 0
for row in range(rows):
for col in range(cols):
_diff += abs(left[row][col] - right[row][col])
return _diff
def backtrack(row: int, col: int, nums_set: int) -> int:
min_cost = rows * cols * max_value
if row == rows:
if validate():
return diff(magick_square, square)
else:
return min_cost
next_col = (col + 1) % cols
next_row = row if next_col > 0 else row + 1
for new_value in range(1, max_value + 1):
if nums_set & (1 << (new_value - 1)):
continue
magick_square[row][col] = new_value
min_cost = min(
min_cost,
backtrack(next_row, next_col, nums_set | (1 << (new_value - 1))),
)
return min_cost
return backtrack(0, 0, 0)
if __name__ == "__main__":
fptr = open(os.environ["OUTPUT_PATH"], "w")
s = []
for _ in range(3):
s.append(list(map(int, input().rstrip().split())))
result = formingMagicSquare(s)
fptr.write(str(result) + "\n")
fptr.close()