425. Word Squares

Given a set of words (without duplicates), find all word squares you can build from them.

A sequence of words forms a valid word square if the kth row and column read the exact same string, where 0 ≤ k< max(numRows, numColumns).

For example, the word sequence ["ball","area","lead","lady"] forms a word square because each word reads the same both horizontally and vertically.

b a l l
a r e a
l e a d
l a d y

Note:

1. There are at least 1 and at most 1000 words.
2. All words will have the exact same length.
3. Word length is at least 1 and at most 5.
4. Each word contains only lowercase English alphabet a-z .

Example 1:

Input:

["area","lead","wall","lady","ball"]


Output:

[
  [ "wall",
    "area",
    "lead",
    "lady"
  ],
  [ "ball",
    "area",
    "lead",
    "lady"
  ]
]


Explanation:

The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).

Example 2:

Input:

["abat","baba","atan","atal"]


Output:

[
  [ "baba",
    "abat",
    "baba",
    "atan"
  ],
  [ "baba",
    "abat",
    "baba",
    "atal"
  ]
]


Explanation:

The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).

Solution

(1) Java

(2) Python

class Solution:
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        def helper():
            if len(path) > 0 and len(path) == len(path[0]):
                rst.append(list(path))
                return
            prefix = ""
            for i in range(len(path)):
                wd = path[i]
                prefix += wd[len(path)]  
            if not prefix in prefixes:
                return
            wds = prefixes[prefix]
            for wd in wds:
                path.append(wd)
                helper()
                path.pop(len(path)-1)                    
        rst = []
        if not words:
            return rst
        prefixes = {}
        length = len(words[0])
        for word in words:
            for i in range(length+1):
                s = word[:i]
                if not s in prefixes:
                    l = [word]
                    prefixes[s] = l
                else:
                    prefixes[s].append(word)
        path = []
        helper()
        return rst

(3) Scala