LeetCode 996 - Number of Squareful Arrays

The problem asks us to count all permutations of a given integer array nums such that the array is squareful, meaning that the sum of every pair of adjacent elements is a perfect square. In other words, for a permutation [a1, a2, a3, ...

LeetCode Problem 996

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Math, Dynamic Programming, Backtracking, Bit Manipulation, Bitmask

Solution

Problem Understanding

The problem asks us to count all permutations of a given integer array nums such that the array is squareful, meaning that the sum of every pair of adjacent elements is a perfect square. In other words, for a permutation [a1, a2, a3, ..., an], each a[i] + a[i+1] must be a perfect square. The array may contain duplicates, and two permutations are considered different if there is any index where the elements differ.

The input is an array of integers, constrained by 1 <= nums.length <= 12 and 0 <= nums[i] <= 10^9. This tells us that brute-force permutations (which would be factorial in length) are feasible only because the array is relatively small, though careful handling of duplicates is required. Large values in nums imply we need to carefully check for perfect squares using arithmetic instead of precomputing all possibilities.

Important edge cases include arrays where all elements are the same, arrays with elements that cannot form any squareful permutation, and arrays with multiple duplicates where permutations might appear identical unless carefully deduplicated.

Approaches

Brute Force

A brute-force approach would generate all permutations of nums and check for each one whether all adjacent sums are perfect squares. This approach works correctly because it exhaustively examines all possibilities, but it is inefficient. The number of permutations is n!, and checking each permutation involves computing the sum of adjacent elements. Furthermore, duplicates would cause overcounting unless we handle them.

Optimal Approach

The key observation is that we do not need to generate all permutations blindly. Instead, we can build the permutation using backtracking, only extending sequences that maintain the squareful property. To efficiently manage duplicates, we sort the array and use a counter to track remaining elements, avoiding duplicate sequences. This approach is essentially a DFS with pruning: we only continue recursion along paths that are valid. Precomputing which elements can follow others as part of a perfect square sum can also optimize the solution.

Approach Time Complexity Space Complexity Notes
Brute Force O(n! * n) O(n) Generates all permutations and checks each
Optimal O(n * 2^n) O(n + n^2) Backtracking with pruning and duplicate handling

Algorithm Walkthrough

  1. Sort and count elements: Sorting helps manage duplicates. Counting each element allows us to efficiently know which elements are available to use at each step.
  2. Build adjacency map: For each pair (x, y) in nums, check if x + y is a perfect square. Store valid next elements in a hash map for fast lookup.
  3. Backtracking function: Recursively try to build sequences by selecting elements that are still available (count > 0) and are valid next elements according to the adjacency map.
  4. Deduplicate using counts: When selecting an element, decrement its count. After recursive calls, restore the count. Skip elements if they are the same as the previous one and the previous one has not been used yet to prevent generating identical sequences.
  5. Base case: When the sequence length equals n, increment the result counter.
  6. Return total: After all recursive calls complete, the counter contains the total number of unique squareful permutations.

Why it works: By only extending valid sequences and using counts to manage duplicates, we ensure that each sequence we count is valid and unique. The backtracking guarantees that all possibilities are explored.

Python Solution

from collections import Counter
from math import isqrt
from typing import List

class Solution:
    def numSquarefulPerms(self, nums: List[int]) -> int:
        def is_square(n: int) -> bool:
            root = isqrt(n)
            return root * root == n
        
        n = len(nums)
        counter = Counter(nums)
        adj = {x: [] for x in counter}
        
        for x in counter:
            for y in counter:
                if is_square(x + y):
                    adj[x].append(y)
        
        def backtrack(x: int, remaining: int) -> int:
            counter[x] -= 1
            if remaining == 1:
                counter[x] += 1
                return 1
            total = 0
            for y in adj[x]:
                if counter[y] > 0:
                    total += backtrack(y, remaining - 1)
            counter[x] += 1
            return total
        
        result = 0
        for x in counter:
            result += backtrack(x, n)
        return result

Implementation walkthrough: The is_square function checks if a number is a perfect square. We build a counter of elements and an adjacency list of valid next elements. The recursive backtrack function builds sequences by decrementing counts, exploring valid continuations, and restoring counts afterward. We start recursion from each unique element and sum all valid sequences.

Go Solution

package main

import (
	"math"
)

func numSquarefulPerms(nums []int) int {
	n := len(nums)
	counter := make(map[int]int)
	for _, x := range nums {
		counter[x]++
	}
	
	adj := make(map[int][]int)
	for x := range counter {
		for y := range counter {
			if isSquare(x + y) {
				adj[x] = append(adj[x], y)
			}
		}
	}
	
	var backtrack func(x, remaining int) int
	backtrack = func(x, remaining int) int {
		counter[x]--
		if remaining == 1 {
			counter[x]++
			return 1
		}
		total := 0
		for _, y := range adj[x] {
			if counter[y] > 0 {
				total += backtrack(y, remaining-1)
			}
		}
		counter[x]++
		return total
	}
	
	result := 0
	for x := range counter {
		result += backtrack(x, n)
	}
	return result
}

func isSquare(n int) bool {
	root := int(math.Sqrt(float64(n)))
	return root*root == n
}

Go-specific notes: We use maps for the counter and adjacency list. Integer arithmetic ensures no overflow because Go ints are large enough for sums of numbers within constraints. Slices are used instead of lists, and recursion is handled similarly to Python.

Worked Examples

Example 1

Input: [1,17,8]

Adjacency map:

1: [8]        (1+8=9)
8: [1,17]     (8+1=9, 8+17=25)
17: [8]       (17+8=25)

Starting from 1: 1 -> 8 -> 17 valid

Starting from 8: 8 -> 1 -> 17 invalid (1+17=18 not square), 8 -> 17 -> 8 invalid (cannot reuse)

Starting from 17: 17 -> 8 -> 1 valid

Output: 2

Example 2

Input: [2,2,2]

Adjacency map: 2: [2] (2+2=4)

Only sequence: [2,2,2]

Output: 1

Complexity Analysis

Measure Complexity Explanation
Time O(n * 2^n) Backtracking explores all valid sequences using remaining counts. Precomputation is O(n^2).
Space O(n^2 + n) Adjacency list uses O(n^2), recursion stack uses O(n), counter uses O(n)

The backtracking approach is efficient because the pruning significantly reduces the factorial blowup, especially with duplicates.

Test Cases

# Provided examples
assert Solution().numSquarefulPerms([1,17,8]) == 2  # basic small input
assert Solution().numSquarefulPerms([2,2,2]) == 1   # all identical elements

# Edge cases
assert Solution().numSquarefulPerms([1]) == 1       # single element
assert Solution().numSquarefulPerms([1,2,3,4]) == 0 # no squareful permutation
assert Solution().numSquarefulPerms([1,8,17,8]) == 4 # duplicates with valid permutations
assert Solution().numSquarefulPerms([0,16,9,7]) == 0 # zeros and mix
Test Why
[1,17,8] Basic example with multiple valid permutations
[2,2,2] Single repeated element
[1] Minimum input size
[1,2,3,4] No possible squareful permutation
[1,8,17,8] Handles duplicates correctly
[0,16,9,7] Mix of zero and large numbers

Edge Cases

All elements identical: Arrays like [2,2,2] only allow one sequence. Failing to deduplicate would overcount. We use counters to handle duplicates.

Single element: [5] should return 1, as a single element trivially satisfies adjacency rules. Base case in recursion ensures this.

No valid sequences: [1,2,3,4] cannot form squareful pairs. Early pruning in backtracking prevents unnecessary exploration and returns