LeetCode #840. Magic Squares In Grid

問題文

原文

A 3 x 3 magic square is a 3 x 3 grid filled with distinct numbers from 1 to 9 such that each row, column, and both diagonals all have the same sum.
Given a row x col grid of integers, how many 3 x 3 magic square subgrids are there?
Note: while a magic square can only contain numbers from 1 to 9, grid may contain numbers up to 15.

日本語訳

3×3 の魔方陣とは、1 から 9 までの 異なる整数を用いて 3×3 のマスを埋め、
各行・各列・2 本の対角線の合計がすべて等しくなるようにしたものを指します。
整数からなる row × col のグリッドが与えられたとき、その中に含まれる 3×3 の魔方陣の部分グリッドの個数を求めてください。
※注意：魔方陣に使われる数は 1〜9 に限られますが、与えられるグリッド自体には 最大 15 までの整数が含まれる場合があります。

つまり，与えられた row x col の grid から，
以下の条件をすべて満たす 3×3 の部分グリッドの個数を数える問題です．

• 3x3 で，1~9 を 1 度ずつ使用
• 各 row，col，対角線の合計が全て同じになる
• 魔法陣の条件は 1~9 だが，grid には 15 までの値が入る可能性がある

アプローチ

単純に計算していたら大変なので，探索の制約を設けるようにします．
それは 魔法陣の中心 (1, 1) は 5 という制約です．
なぜこれが成り立つか見ていきましょう．

なぜ魔法陣の中心は5？

計算していけば割と簡単です．

1. 魔法陣の総和と各行，列，対角線の和
   まず，魔法陣の総和は 45 になります．
   このとき，各行・列・対角線の和は同一値 15 に固定されます．
   例えば，それぞれの行の和を $S$ とした場合，

   \begin{align}
   1 + 2 + 3 + ... + 9 &= 45 \\
   3S &= 45 \\
   S &= 15
   \end{align}
   

   という形でそれぞれの行の和 $S$ 求められますね．
2. 中心が 5 となる理由
   3x3 のグリッドの各値を下記のように $a〜i$ とする．

   \begin{matrix}
   a & b & c \\
   d & e & f \\
   g & h & i
   \end{matrix}
   

   中心 $e$ を通る線が 4 本あり，それぞれ総和は 15 なので，下記のように表すことができます．

   \begin{align}
   d + e + f &= 15 \\
   b + e + h &= 15 \\
   a + e + i &= 15 \\
   c + e + g &= 15 \\
   
   \end{align}
   

   行・列・対角線 4 つ分であり，そのほかの 8 要素は "全体の総和 - $e$"になるので，

   \begin{align}
   (d+b+a+c+f+h+i+g)+4e &=60 \\
   d+b+a+c+f+h+i+g &= 45 -e \\
   3e &= 15 \\
   e &= 5
   \end{align}
   

よって，中心 $e$ は 5 に固定される．

こうなると，探索範囲は一気に限られます．
あとは下記のような条件を満たすように探せば OK です．

• 3x3 のグリッド内には 1~9 の値がユニークに入る
• 各行・列・対角線は総和が 15 になる

Python での実装例

class Solution:
    def numMagicSquaresInside(self, grid: List[List[int]]) -> int:
        ans = 0
        n = len(grid)
        m = len(grid[0])
        lines = [
            (0,1,2),(3,4,5),(6,7,8),
            (0,3,6),(1,4,7),(2,5,8),
            (0,4,8),(2,4,6),
        ]
        for r in range(1, n-1):
            for c in range(1, m-1):
                if grid[r][c] != 5:
                    continue
                subgrid = [
                    grid[r-1][c-1], grid[r-1][c], grid[r-1][c+1],
                    grid[r][c-1], grid[r][c], grid[r][c+1],
                    grid[r+1][c-1], grid[r+1][c], grid[r+1][c+1],
                ]
                maxv = max(subgrid)
                if maxv > 9:
                    continue
                if len(set(subgrid)) != 9:
                    continue
                if any(subgrid[i]+subgrid[j]+subgrid[k] != 15 for i,j,k in lines):
                    continue
                ans += 1
        return ans

さいごに

今回は気付けば簡単ですが，気づかないと沼りそうな問題でした．
こういうのがあるから面白いですね！