840. Magic Squares In Grid

https://leetcode.com/problems/magic-squares-in-grid/

JavaScript

/**
 * @param {number[][]} grid
 * @return {number}
 */
var numMagicSquaresInside = function(grid) {
  const rows = grid.length
  const cols = grid[0].length

  if (rows < 3 || cols < 3) {
    return 0
  }

  let count = 0
  for (let i = 0; i < rows - 2; i++) {
    for (let j = 0; j < cols - 2; j++) {
      const subgrid = grid.slice(i, i + 3).map(row => row.slice(j, j + 3))

      if (subgrid[1][1] !== 5) {
        continue
      }

      const nums = subgrid[0].concat(subgrid[1]).concat(subgrid[2])
      const numMatch = new Set(nums).size === 9 && nums.every(n => n >= 1 && n <= 9)
      if (!numMatch) {
        continue
      }

      const rowSums = subgrid.map(r => r[0] + r[1] + r[2])
      const colSums = [0, 1, 2].map(c => subgrid[0][c] + subgrid[1][c] + subgrid[2][c])
      const sumD1 = subgrid[0][0] + subgrid[1][1] + subgrid[2][2]
      const sumD2 = subgrid[0][2] + subgrid[1][1] + subgrid[2][0]

      if (new Set(rowSums.concat(colSums).concat([sumD1, sumD2])).size === 1) {
        count += 1
      }
    }
  }

  return count
}

Python

class Solution:

    def numMagicSquaresInside(self, grid):
        """
        :type grid: List[List[int]]
        :rtype: int
        """
        rows = len(grid)
        cols = len(grid[0])

        if rows < 3 or cols < 3:
            return 0

        count = 0

        for i in range(rows - 2):
            for j in range(cols - 2):
                subgrid = [row[j:j + 3] for row in grid[i:i + 3]]

                if subgrid[1][1] != 5:
                    continue

                num_match = set([x for row in subgrid for x in row]) == set(
                    range(1, 10))
                if not num_match:
                    continue

                row_sums = [sum(row) for row in subgrid]
                col_sums = [sum([subgrid[r][c] for r in range(3)])
                            for c in range(3)]
                sum_d1 = subgrid[0][0] + subgrid[1][1] + subgrid[2][2]
                sum_d2 = subgrid[0][2] + subgrid[1][1] + subgrid[2][0]

                if len(set(row_sums + col_sums + [sum_d1, sum_d2])) == 1:
                    count += 1

        return count