Synonymous Sentences Python Solution

lcpy gen -n 1258   # by problem number
lcpy gen -s synonymous_sentences   # by problem name

Input: synonyms = [["happy","joy"],["sad","sorrow"],["joy","cheerful"]], text = "I am happy today but was sad yesterday"
Output: ["I am cheerful today but was sad yesterday","I am cheerful today but was sorrow yesterday","I am happy today but was sad yesterday","I am happy today but was sorrow yesterday","I am joy today but was sad yesterday","I am joy today but was sorrow yesterday"]
Explanation: From the synonyms, happy, joy and cheerful are equivalent, and sad and sorrow are equivalent. Replacing each synonymous word independently gives 2 * 3 = 6 sentences.

Input: synonyms = [["happy","joy"],["cheerful","glad"]], text = "I am happy today but was sad yesterday"
Output: ["I am happy today but was sad yesterday","I am joy today but was sad yesterday"]

class UnionFind:
    def __init__(self, n: int):
        self.parent: list[int] = list(range(n))
        self.size: list[int] = [1] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a: int, b: int) -> None:
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]


class Solution:
    # Time: O((s + w) * n)
    # Space: O(s + w)
    def generate_sentences(self, synonyms: list[list[str]], text: str) -> list[str]:
        words = sorted({word for pair in synonyms for word in pair})
        index = {word: i for i, word in enumerate(words)}
        uf = UnionFind(len(words))
        for first, second in synonyms:
            uf.union(index[first], index[second])

        groups: dict[int, list[str]] = {}
        for word in words:
            groups.setdefault(uf.find(index[word]), []).append(word)

        sentence = text.split()
        result: list[str] = []
        current: list[str] = []

        def dfs(i: int) -> None:
            if i == len(sentence):
                result.append(" ".join(current))
                return
            word = sentence[i]
            if word in index:
                for alt in groups[uf.find(index[word])]:
                    current.append(alt)
                    dfs(i + 1)
                    current.pop()
            else:
                current.append(word)
                dfs(i + 1)
                current.pop()

        dfs(0)
        return sorted(result)