magic-sq-constraint.py

#!/usr/bin/env python3
from constraint import *
N, magic = 3, 15
# N, magic = 4, 34

sq = N ** 2


def check(*a):
    for j in range(N):
        if sum(a[N * j:N * (j + 1)]) != magic:
            return False

    # 147 258 369
    for i in range(N):
        s = 0
        for j in range(N):
            s += a[i + (j * N)]
        if s != magic:
            return False

    # 159 | 1,6,11,16
    s = 0
    for i in range(N):
        s += a[(N + 1) * i]
    if s != magic:
        return False

    # 357[246] | 4,7,10,13[3,6,9,12]
    s = 0
    for i in range(1, N + 1):
        s += a[(N - 1) * i]
    if s != magic:
        return False
        
    return True


def main():
    p = Problem()
    p.addVariables(range(sq), range(1, sq + 1))
    p.addConstraint(AllDifferentConstraint())
    p.addConstraint(ExactSumConstraint((sq * (sq + 1)) / 2))
    p.addConstraint(check)
    print(len(p.getSolutions()))


if __name__ == '__main__':
    main()

magic-sq.py

#!/usr/bin/env python3
cnt = 0
def print_square(a, N):
    for i in range(N):
        print(a[N * i:N * (i + 1)])
    print("___")

def check(a, N, magic):
    # Checked prior
    # for j in range(N):
    #     if sum(a[N * j:N * (j + 1)]) != magic:
    #         return False

    # 147 258 369
    for i in range(N):
        s = 0
        for j in range(N):
            s += a[i + (j * N)]
        if s != magic:
            return False

    # 159 | 1,6,11,16
    s = 0
    for i in range(N):
        s += a[(N + 1) * i]
    if s != magic:
        return False
   
    # 357[246] | 4,7,10,13[3,6,9,12]
    s = 0
    for i in range(1, N + 1):
        s += a[(N - 1) * i]
    if s != magic:
        return False

    return True

def permute(a, i, N, magic):
    global cnt
    sq = N ** 2
    if i == sq:
        if check(a, N, magic):
            cnt += 1
            print_square(a, N)
        return
    for j in range(N):
        if i == N * (j + 1) and sum(a[N * j:N * (j + 1)]) != magic:
            return
    exist_map = [False] * sq
    for j in range(i):
        exist_map[a[j] - 1] = True
    for e in range(1, sq + 1):
        if exist_map[e - 1]:
            continue
        a[i] = e
        permute(a, i + 1, N, magic)
        # print(a, i, exist_map)


def main():
    n, magic = 3, 15
    # n, magic = 4, 34
    permute(list(range(1, (n ** 2) + 1)), 0, n, magic)
    print('Count:',cnt)


if __name__ == '__main__':
    main()