fork(3) download
  1. from collections import defaultdict, deque
  2. from copy import copy, deepcopy
  3.  
  4. def magicSum(n):
  5. return int((n*n * (n*n + 1)) / 6)
  6.  
  7. def validate(sumDict, n):
  8. for k, v in sumDict.items():
  9. if v > magicSum(n):
  10. return False
  11. return True
  12.  
  13. def check(sumDict, n):
  14. for k, v in sumDict.items():
  15. if v != magicSum(n):
  16. return False
  17. return True
  18.  
  19. def isValid(m, n):
  20. rowSum = defaultdict(int)
  21. colSum = defaultdict(int)
  22. diagSum = defaultdict(int)
  23.  
  24. isLeft = False
  25.  
  26. for i in range(n):
  27. for j in range(n):
  28. if m[i][j] == 0: isLeft = True
  29. rowSum[i] += m[i][j]
  30. colSum[j] += m[i][j]
  31. if i == j: diagSum[0] += m[i][j]
  32. if i + j == n - 1: diagSum[n - 1] += m[i][j]
  33.  
  34. if isLeft:
  35. return (validate(rowSum, n) and validate(colSum, n) and validate(diagSum, n))
  36. return (check(rowSum, n) and check(colSum, n) and check(diagSum, n))
  37.  
  38. def next(cur, m, n):
  39. possible = []
  40. for i in range(n):
  41. for j in range(n):
  42. if m[i][j] == 0:
  43. nextM = deepcopy(m)
  44. nextM[i][j] = cur
  45. if isValid(nextM, n):
  46. possible.append(nextM)
  47. return possible
  48.  
  49. def printM(m):
  50. for i in range(len(m)):
  51. print(m[i])
  52. print("\n\n")
  53.  
  54. def gen(n):
  55. startM = [[0 for x in range(n)] for y in range(n)]
  56. magic = []
  57. Q = deque([(1, startM)])
  58. while len(Q):
  59. state = Q.popleft()
  60. cur = state[0]
  61. m = state[1]
  62. if cur == n * n + 1:
  63. magic.append(m)
  64. printM(m)
  65. continue
  66. for w in next(cur, m, n):
  67. Q.append((cur + 1, w))
  68. return magic
  69.  
  70. magic = gen(3)
Time limit exceeded #stdin #stdout 5s 76480KB
stdin
Standard input is empty
stdout
Standard output is empty