Skip to content

996. Number of Squareful Arrays 👍

  • Time: $O(n \cdot 2^n)$
  • Space: $O(n \cdot 2^n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
 public:
  int numSquarefulPerms(vector<int>& nums) {
    int ans = 0;
    ranges::sort(nums);
    dfs(nums, vector<bool>(nums.size()), {}, ans);
    return ans;
  }

 private:
  void dfs(vector<int>& nums, vector<bool>&& used, vector<int>&& path,
           int& ans) {
    if (path.size() > 1 && !isSquare(path.back() + path[path.size() - 2]))
      return;
    if (path.size() == nums.size()) {
      ++ans;
      return;
    }

    for (int i = 0; i < nums.size(); ++i) {
      if (used[i])
        continue;
      if (i > 0 && nums[i] == nums[i - 1] && !used[i - 1])
        continue;
      used[i] = true;
      path.push_back(nums[i]);
      dfs(nums, std::move(used), std::move(path), ans);
      path.pop_back();
      used[i] = false;
    }
  }

  bool isSquare(int num) {
    const int root = sqrt(num);
    return root * root == num;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
  public int numSquarefulPerms(int[] nums) {
    boolean[] used = new boolean[nums.length];
    Arrays.sort(nums);
    dfs(nums, used, new ArrayList<>());
    return ans;
  }

  private int ans = 0;

  private void dfs(int[] nums, boolean[] used, List<Integer> path) {
    if (path.size() > 1 && !isSquare(path.get(path.size() - 1) + path.get(path.size() - 2)))
      return;
    if (path.size() == nums.length) {
      ++ans;
      return;
    }

    for (int i = 0; i < nums.length; ++i) {
      if (used[i])
        continue;
      if (i > 0 && nums[i] == nums[i - 1] && !used[i - 1])
        continue;
      used[i] = true;
      path.add(nums[i]);
      dfs(nums, used, path);
      path.remove(path.size() - 1);
      used[i] = false;
    }
  }

  private boolean isSquare(int num) {
    int root = (int) Math.sqrt(num);
    return root * root == num;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution:
  def numSquarefulPerms(self, nums: list[int]) -> int:
    ans = 0
    used = [False] * len(nums)

    def isSquare(num: int) -> bool:
      root = math.isqrt(num)
      return root * root == num

    def dfs(path: list[int]) -> None:
      nonlocal ans
      if len(path) > 1 and not isSquare(path[-1] + path[-2]):
        return
      if len(path) == len(nums):
        ans += 1
        return

      for i, a in enumerate(nums):
        if used[i]:
          continue
        if i > 0 and nums[i] == nums[i - 1] and not used[i - 1]:
          continue
        used[i] = True
        dfs(path + [a])
        used[i] = False

    nums.sort()
    dfs([])
    return ans