Transitive synonym groups from word contexts

Question: Transitive synonym groups from word contexts

Companies: Glean

Difficulty: Medium

Problem statement

Group words into synonym sets where sharing a four-word context links them transitively. Reported on 1point3acres' Glean bank; reconstructed in our own words.

Clarify the setup

A search engine wants to learn synonyms from text: words that appear in exactly the same surroundings are probably interchangeable ("I bought a car from the dealer" and "I bought a vehicle from the dealer"). Precisely:

• The context of a word occurrence is the two words before it and the two words after it: four words.
• Two different words are directly linked if some context occurs with both.
• Synonym groups are the transitive closure: if car ~ vehicle and vehicle ~ automobile, all three are one group.

Return the groups with at least two words. Ask:

• Edges of a sentence. A word with fewer than two words on either side has no full context; skip it (or pad with sentence markers; ask).
• Normalization. Lowercase and strip punctuation first.
• Output order. Sort words within a group and groups by their first word, so results are testable.
• Scale. Millions of sentences: the solution should be close to linear in the input.

Approach

Brute force: for every pair of distinct words, check whether they share a context: O(V² × contexts). Too slow.

Bucket by context, then union-find.

1. Slide a 5-word window over each sentence. The middle word is the candidate, the other four (as a tuple) are the key.
2. Keep first_word_for[context]. The first time a context appears, remember its word. Each later occurrence with a different word is a link: union(word, first_word_for[context]). Linking every word to the first one in the bucket is enough to connect the whole bucket; no need for all pairs.
3. Union-find (disjoint sets) with path compression and union by size merges groups in nearly O(1) amortized per operation.
4. At the end, group words by their root and keep groups of two or more.

Every window is processed once with O(1) dictionary work, so the whole thing is O(N) for N words of text (times the inverse-Ackermann factor of union-find, which is at most 4 for any real input).

Why transitive closure is risky. One bad link merges two large groups: "a bank of the river" and "a side of the …" might share a context with unrelated words, and transitivity spreads the error. Follow-ups cover requiring several shared contexts and ignoring very common contexts.

Code

import re
from collections import defaultdict
from typing import Dict, Iterable, List, Tuple

WORD = re.compile(r"[a-z0-9']+")


class DisjointSet:
    def __init__(self) -> None:
        self.parent: Dict[str, str] = {}
        self.size: Dict[str, int] = {}

    def find(self, x: str) -> str:
        self.parent.setdefault(x, x)
        self.size.setdefault(x, 1)
        root = x
        while self.parent[root] != root:
            root = self.parent[root]
        while self.parent[x] != root:             # path compression
            self.parent[x], x = root, self.parent[x]
        return root

    def union(self, a: str, b: str) -> None:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return
        if self.size[ra] < self.size[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        self.size[ra] += self.size[rb]


def tokenize(sentence: str) -> List[str]:
    return WORD.findall(sentence.lower())


def synonym_groups(sentences: Iterable[str], min_shared: int = 1) -> List[List[str]]:
    """Groups of words linked by shared 4-word contexts, transitively.

    min_shared > 1 (part 2) only links two words that share that many
    distinct contexts, which keeps one accidental match from merging groups.
    """
    first_word_for: Dict[Tuple[str, ...], str] = {}
    shared: Dict[Tuple[str, str], set] = defaultdict(set)
    ds = DisjointSet()
    for sentence in sentences:
        words = tokenize(sentence)
        for i in range(2, len(words) - 2):
            word = words[i]
            context = (words[i - 2], words[i - 1], words[i + 1], words[i + 2])
            ds.find(word)
            first = first_word_for.setdefault(context, word)
            if first == word:
                continue
            pair = (min(first, word), max(first, word))
            shared[pair].add(context)
            if len(shared[pair]) >= min_shared:
                ds.union(first, word)
    groups: Dict[str, List[str]] = defaultdict(list)
    for word in ds.parent:
        groups[ds.find(word)].append(word)
    return sorted(sorted(g) for g in groups.values() if len(g) > 1)

With min_shared = 1 the shared map is not needed; it is there for part 2. To keep memory down at scale, store only pairs below the threshold and drop a pair's set once it is linked.

Complexity

• Time: O(N α(V)) for N words of text and V distinct words; effectively linear.
• Space: O(number of distinct contexts) for first_word_for, plus O(V) for union-find. Distinct 4-word contexts can number in the hundreds of millions on web-scale text, which is why the follow-ups move this to a batch job.

Test cases

def test_direct_and_transitive_links():
    sentences = [
        "i bought a car from the dealer",
        "i bought a vehicle from the dealer",
        "she parked the vehicle near the house",
        "she parked the automobile near the house",
        "the cat sat on the mat",
    ]
    assert synonym_groups(sentences) == [["automobile", "car", "vehicle"]]


def test_normalization_and_short_sentences():
    sentences = ["Oh, we LOVE Dogs, really much!", "oh we love cats really much", "too short"]
    assert synonym_groups(sentences) == [["cats", "dogs"]]
    assert synonym_groups([]) == []


def test_edges_have_no_full_context():
    # "x" and "y" are at the sentence start, so they have no 4-word context
    assert synonym_groups(["x a b c d", "y a b c d"]) == []


def test_two_separate_groups():
    sentences = [
        "go to the shop to buy milk", "go to the store to buy milk",
        "a very big house on hill", "a very large house on hill",
    ]
    assert synonym_groups(sentences) == [["big", "large"], ["shop", "store"]]

def test_min_shared_blocks_a_single_accidental_link():
    sentences = [
        "we sat by the bank of the river",
        "we sat by the side of the river",
        "money in the bank is safe now",
        "money in the vault is safe now",
        "funds in the bank are low today",
        "funds in the vault are low today",
    ]
    assert synonym_groups(sentences) == [["bank", "side", "vault"]]
    assert synonym_groups(sentences, min_shared=2) == [["bank", "vault"]]


def test_disjoint_set_compresses_paths():
    ds = DisjointSet()
    for a, b in [("a", "b"), ("c", "d"), ("b", "d"), ("e", "a")]:
        ds.union(a, b)
    root = ds.find("e")
    assert all(ds.find(x) == root for x in "abcde")
    assert ds.size[root] == 5

Follow-ups

1. Common contexts are noise. Contexts like ("of", "the", "and", "a") match thousands of words. Ignore contexts whose bucket has more than, say, 50 distinct words, or weight links by how rare the context is.
2. Scale out. As a MapReduce or Spark job: map each window to (context, word), group by context, emit (word, first_word) links for small buckets, then run connected components on the link graph (iterative label propagation or a union-find over a sorted edge list).
3. Softer similarity. Exact four-word matches are rare for rare words. Use partial contexts (one or two neighbors), count vectors, or embeddings, and link words whose similarity passes a threshold. Transitive closure becomes clustering with a minimum similarity.
4. Using it in search. Synonyms expand queries (car also matches automobile), but precision suffers if groups are wrong. Use groups as a ranking signal or require query-log evidence (users who search one also click results for the other) before expanding.
5. Antonyms. "The price went up today" and "the price went down today" share contexts, so distributional methods link antonyms. Filter with a lexical resource or click data.

Interviewer expectations

Mid-level (L4)

• Builds contexts with a sliding window and links words that share one.
• Uses union-find or graph BFS for the transitive closure and returns sorted groups.
• Handles short sentences and normalization.

Senior (L5)

• Avoids all-pairs work by linking each word to the first word of its context bucket, and states the near-linear cost.
• Implements union-find with path compression and union by size.
• Points out how transitive closure spreads one bad link, and adds a minimum number of shared contexts.

Staff (L6+)

• Filters frequent contexts, and moves the computation to a batch job with connected components at web scale.
• Discusses softer similarity (embeddings, partial contexts) and the precision risk of using synonyms in search, including antonyms.
• Proposes how to validate groups with query logs before they affect ranking.