Chu-Liu-Edmonds, in depth: decoding dependency trees from a score matrix, the single-root constraint, and testing against brute force

Why greedy heads are not enough

The obvious decoder takes, for every word, the head with the highest score. Each word then has one head, but nothing stops two words choosing each other, leaving a cycle disconnected from ROOT. The obvious repair, deleting the weakest arc in the cycle and attaching somewhere else, is not optimal either: the right arc to remove depends on which outside arc can best replace it, and on what the rest of the tree looks like after that.

Undirected minimum spanning tree algorithms such as Kruskal's do not help, because direction matters: an arc from read to she is a different decision from an arc from she to read, and the constraint is on in-degree, not on connectivity alone. Chu-Liu-Edmonds resolves cycles exactly by contracting them and solving a smaller problem.

The algorithm: greedy, contract, expand

The algorithm has three steps, applied recursively.

1. Greedy heads. For every node except ROOT, pick the incoming arc with the highest score. If these arcs form no cycle, they are the optimal tree; stop.
2. Contract. Otherwise take one cycle C and replace its nodes with a single node c. An arc from an outside node u into cycle node v gets the score S[u][v] minus S[head(v)][v]: entering the cycle at v means giving up v's cycle arc, so the arc is worth only its gain over what it displaces. Keep the best such arc per outside node. An arc leaving the cycle from any member w to an outside node keeps its score; keep the best per target.
3. Recurse and expand. Solve the smaller graph. The arc it picks into c names the cycle node v where the tree enters; v takes that outside head, every other cycle node keeps its cycle arc, and arcs out of c are mapped back to the member they came from.

Why is this optimal? Every arborescence must give each cycle node exactly one head, and at least one cycle node must take its head from outside the cycle, otherwise the cycle is unreachable from ROOT. An optimal tree can always be chosen that keeps all but one cycle arc, so the cost of the cycle is fixed except for the single point of entry, and the relative scores price exactly that choice. Each contraction removes at least one node, so there are at most n levels of recursion.

A dense-matrix implementation

The implementation works on a list-of-lists matrix with NEG = float("-inf") for impossible arcs, so you can pass a parser's score matrix after masking the diagonal and the ROOT column. Cycle detection walks head pointers with a three-colour mark, a simpler cousin of the depth-first search used for Tarjan's strongly connected components.

NEG = float("-inf")

def find_cycle(head):
    color = [0] * len(head)          # 0 new, 1 on current walk, 2 finished
    for start in range(len(head)):
        path, v = [], start
        while v != -1 and color[v] == 0:
            color[v] = 1
            path.append(v)
            v = head[v]
        if v != -1 and color[v] == 1:
            return path[path.index(v):]
        for u in path:
            color[u] = 2
    return None

def max_arborescence(S, root=0):
    """S[h][d] scores arc h -> d. Returns head list with head[root] = -1."""
    n = len(S)
    head = [-1] * n
    for d in range(n):
        if d == root:
            continue
        head[d] = max((h for h in range(n) if h != d), key=lambda h: S[h][d])
        if S[head[d]][d] == NEG:
            raise ValueError(f"node {d} has no incoming arc")
    cycle = find_cycle(head)
    if cycle is None:
        return head
    in_c = set(cycle)
    keep = [v for v in range(n) if v not in in_c]
    idx = {v: i for i, v in enumerate(keep)}
    c = len(keep)                                  # index of the contracted node
    T = [[NEG] * (c + 1) for _ in range(c + 1)]
    enter, leave = {}, {}
    for u in keep:
        for v in keep:
            if u != v:
                T[idx[u]][idx[v]] = S[u][v]
        best, arg = NEG, None                      # u -> cycle, priced by what it displaces
        for v in cycle:
            if S[u][v] != NEG and S[u][v] - S[head[v]][v] > best:
                best, arg = S[u][v] - S[head[v]][v], v
        T[idx[u]][c], enter[u] = best, arg
        best, arg = NEG, None                      # cycle -> u, best member
        for w in cycle:
            if S[w][u] > best:
                best, arg = S[w][u], w
        T[c][idx[u]], leave[u] = best, arg
    sub = max_arborescence(T, idx[root])
    new_head = list(head)                          # cycle nodes keep cycle arcs by default
    for v in keep:
        h = sub[idx[v]]
        new_head[v] = -1 if v == root else (leave[v] if h == c else keep[h])
    u = keep[sub[c]]                               # outside node that wins entry
    new_head[enter[u]] = u                         # breaks the cycle where it enters
    return new_head

def score(S, head):
    return sum(S[h][d] for d, h in enumerate(head) if h != -1)

Each level does O(n²) work to build the contracted matrix, and there are at most n levels, so this version is O(n³) in the worst case. Tarjan's 1977 refinement reaches O(n²) on dense graphs by never rebuilding the matrix, using union-find to track which original nodes have merged. For sentences under a hundred words the simple version runs in milliseconds in Python; profile before optimising.

Worked example: ROOT she read books today

Score the sentence ROOT she read books today with the arcs below (all others are low). The greedy step picks read as head of she (15), she as head of read (14), read as head of books (18) and read as head of today (9). she and read point at each other: a cycle with total score 29.

Contract {she, read} into C. Entry arcs are priced by what they displace: ROOT to read scores 12 - 14 = -2, ROOT to she scores 4 - 15 = -11, books to read scores 6 - 14 = -8, today to read scores 1 - 14 = -13. The best entry for C is from ROOT at -2. Arcs out of C keep their best member: C to books is 18 (from read) and C to today is 9.

In the contracted graph the greedy heads are ROOT for C, C for books and C for today. There is no cycle, so expand. ROOT enters the cycle at read, so read's cycle arc from she is dropped and she keeps read as its head. The final tree is ROOT to read, read to she, read to books, read to today, with score 12 + 15 + 18 + 9 = 54, which brute force confirms is the maximum.

W = ["ROOT", "she", "read", "books", "today"]
S = [[NEG] * 5 for _ in range(5)]
for (h, d), s in {(0, 1): 4, (0, 2): 12, (0, 3): 2, (0, 4): 3, (1, 2): 14, (2, 1): 15,
                  (2, 3): 18, (2, 4): 9, (3, 2): 6, (3, 4): 2, (4, 2): 1, (1, 3): 3,
                  (3, 1): 2, (1, 4): 1, (4, 1): 1, (4, 3): 1}.items():
    S[h][d] = s
head = max_arborescence(S)
print([(W[h], W[d]) for d, h in enumerate(head) if h != -1], score(S, head))
# [('read', 'she'), ('ROOT', 'read'), ('read', 'books'), ('read', 'today')] 54

The single-root constraint

Universal Dependencies and most other treebanks require exactly one word to attach to ROOT. Chu-Liu-Edmonds knows nothing about this: if two words both score well as ROOT's children, the optimal arborescence happily gives ROOT two children. Take three words where ROOT to word 1 scores 10, ROOT to word 2 scores 9, and the best alternative head for word 2 scores only 3. The unconstrained optimum attaches both words to ROOT, plus 1 to 3 at 8, for 27. The best single-root tree is ROOT to 1, 1 to 2, 1 to 3, for 21.

Picking the best-scoring ROOT child first and then decoding is a common heuristic, and it is wrong: it can lose to a tree hanging from a different word. Two exact fixes exist. The first tries every candidate child c, forbids every other ROOT arc, decodes, and keeps the best; it costs n decodes. The second needs one decode: subtract a constant M from every ROOT arc, where M exceeds the spread of all possible tree scores. Every extra ROOT child then costs more than any rearrangement can gain, so the optimum uses exactly one ROOT child whenever such a tree exists, and among those it is the best by the original scores.

def single_root(S, root=0):
    M = 1 + sum(abs(x) for row in S for x in row if x != NEG)
    T = [row[:] for row in S]
    for d in range(len(S)):
        if d != root and T[root][d] != NEG:
            T[root][d] -= M                    # each ROOT child now costs M
    return max_arborescence(T, root)           # score it with the original S

The penalty must be exact to be safe. With integer or rational scores it is; with float32 log-probabilities a huge M swallows the small differences between candidate trees, so either scale scores to integers or fall back to the n-decode loop. Zmigrod, Vieira and Cotterell (2020) studied this constraint and gave an O(n²) algorithm for it, adapting work of Gabow and Tarjan; whichever method you ship, test it against brute force.

Testing against brute force

Decoders are easy to get subtly wrong, and wrong ones still return valid-looking trees. Test against exhaustive search on small random matrices, and check both the unconstrained and single-root versions.

import itertools, random

def is_tree(head, root=0):
    for v in range(len(head)):
        seen, u = set(), v
        while u != root:
            if u == -1 or u in seen:
                return False
            seen.add(u)
            u = head[u]
    return True

def brute(S, root=0, single=False):
    n, best = len(S), NEG
    others = [v for v in range(n) if v != root]
    for choice in itertools.product(range(n), repeat=len(others)):
        head = [-1] * n
        for v, h in zip(others, choice):
            head[v] = h
        if any(h == v or S[h][v] == NEG for v, h in zip(others, choice)):
            continue
        if not is_tree(head, root):
            continue
        if single and sum(head[v] == root for v in others) != 1:
            continue
        best = max(best, score(S, head))
    return best

random.seed(7)
for _ in range(3000):
    n = random.randint(2, 6)
    S = [[NEG if h == d or d == 0 else random.randint(-5, 9) for d in range(n)]
         for h in range(n)]
    assert score(S, max_arborescence(S)) == brute(S)
    assert score(S, single_root(S)) == brute(S, single=True)

Small integer scores create many ties, which is deliberate: compare scores, not arc sets, because tied optima can differ. Six nodes means at most 6⁵ = 7,776 head assignments per case, so the test runs in seconds.

Using it in a parser

In a neural parser, the network produces arc scores on the GPU and decoding runs on the CPU, one sentence at a time, because the recursion is sequential and branchy. Three practical rules follow.

• Use log-probabilities. The tree score is a sum, so feed log-softmax outputs, not probabilities, or the decoder maximises the wrong quantity.
• Mask before decoding. Set the diagonal and the ROOT column to NEG. A NaN from a bad batch poisons every max comparison silently; assert the matrix is finite or NEG first.
• Decode only at inference. Training usually optimises per-word head classification or a tree-level loss computed with the matrix-tree theorem; exact decoding is needed only when you must output a well-formed tree.

If your treebank is projective, meaning arcs never cross when drawn above the sentence, Eisner's O(n³) dynamic programme finds the best projective tree and is the better choice. Chu-Liu-Edmonds finds the best tree of any shape, which matters for languages with freer word order where crossing arcs are common.

Failure modes

What to do next

1. Copy max_arborescence, single_root and the brute-force test, and run the test until it passes 3,000 random cases.
2. Trace the five-node example by hand and confirm the contracted scores -2, -11, -8 and -13.
3. Decode a real parser's score matrix with and without the single-root constraint and count how often they differ.
4. If your data is projective, implement Eisner's algorithm and compare accuracy and speed.
5. If decoding shows up in profiles, implement the O(n²) variant and keep the simple version as a test oracle.
6. Add the matrix masking and finite-value assertions to your inference path.