Leetcode: 425. Word Squares

Mistake: Incorrect evaluation of solution's viability. Spent most of the time trying to reduce the complexity, when the problem required a Backtracking.
Mistake: High constant factor in the implementation. First Backtracking wasn't optimized enough.
Can we use brute-force to solve the problem? Generate all possible squares while selecting the next word that match the current prefix so far.

from typing import List


class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:
        N = len(words)
        M = len(words[0])

        if M == 1:
            return list(set(words))

        pref = {}
        for w in words:
            for i in range(M):
                p = w[0 : i + 1]
                pref.setdefault(p, set())
                pref[p].add(w)

        ans = []

        def bt(i, cur):
            if i == M:
                ans.append(cur[:])
                return

            for w in pref.get("".join(v[i] for v in cur), []):
                cur.append(w)
                bt(i + 1, cur)
                cur.pop()

        for w in words:
            bt(1, [w])

        return ans


assert Solution().wordSquares(["area", "lead", "wall", "lady", "ball"]) == [
    ["ball", "area", "lead", "lady"],
    ["wall", "area", "lead", "lady"],
]
assert Solution().wordSquares(["abat", "baba", "atan", "atal"]) == [
    ["baba", "abat", "baba", "atal"],
    ["baba", "abat", "baba", "atan"],
]