Phylogeny

7 problems from Rosalind — Bioinformatics Textbook Track. Press Run on any block to execute it in your browser.

# Rosalind: BA7A — Compute Distances Between Leaves
# https://rosalind.info/problems/ba7a/
#
# Given: An integer n and the adjacency list of a weighted tree with n leaves.
# Return: The n x n matrix of path lengths between leaves.

let leaf_count = 4
let tree = {
    "0": [{ to_node: "4", weight: 11 }],
    "1": [{ to_node: "4", weight: 2 }],
    "2": [{ to_node: "5", weight: 6 }],
    "3": [{ to_node: "5", weight: 7 }],
    "4": [{ to_node: "0", weight: 11 }, { to_node: "1", weight: 2 }, { to_node: "5", weight: 4 }],
    "5": [{ to_node: "4", weight: 4 }, { to_node: "3", weight: 7 }, { to_node: "2", weight: 6 }],
}

# In a tree there is exactly one path between any two nodes, so there is nothing
# to optimise — no Dijkstra, no relaxation. Walking outwards from each leaf and
# recording what it costs to arrive is enough, and each node is reached once.
# That uniqueness is what a tree buys, and it is also why these distances can be
# inverted to recover the tree in BA7C.
fn distances_from(start, graph) {
    let seen = { }
    seen[start] = 0
    let frontier = [start]
    while len(frontier) > 0 {
        let node = frontier[0]
        frontier = slice(frontier, 1, len(frontier))
        for edge in graph[node] {
            if contains(keys(seen), edge.to_node) == false {
                seen[edge.to_node] = seen[node] + edge.weight
                frontier = push(frontier, edge.to_node)
            }
        }
    }
    seen
}

let leaf_distances = range(0, leaf_count) |> map(|i| {
    let reach = distances_from(str(i), tree)
    range(0, leaf_count) |> map(|j| reach[str(j)])
})

println("Result:")
for line in leaf_distances { println("  " + (line |> map(|d| str(d)) |> join("\t"))) }
println("Expected:")
println("  0 13 21 22 / 13 0 12 13 / 21 12 0 13 / 22 13 13 0")

fn test_ba7a_distances_between_leaves() {
    let expected = [[0, 13, 21, 22], [13, 0, 12, 13], [21, 12, 0, 13], [22, 13, 13, 0]]
    assert leaf_distances == expected, "BA7A: got " + str(leaf_distances)
    # A distance matrix from a tree is symmetric, zero on the diagonal, and obeys
    # the triangle inequality — properties BA7C relies on to rebuild the tree.
    for i in range(0, leaf_count) {
        assert leaf_distances[i][i] == 0, "BA7A: a leaf is no distance from itself"
        for j in range(0, leaf_count) {
            assert leaf_distances[i][j] == leaf_distances[j][i], "BA7A: the leaf_distances must be symmetric"
            for k in range(0, leaf_count) {
                assert leaf_distances[i][j] <= leaf_distances[i][k] + leaf_distances[k][j],
                    "BA7A: the triangle inequality fails at "
                        + str(i) + "," + str(j) + "," + str(k)
            }
        }
    }
}

# Rosalind: BA7B — Compute Limb Length
# https://rosalind.info/problems/ba7b/
#
# Given: An integer n, a leaf j, and an additive n x n distance matrix.
# Return: The length of the limb connecting leaf j to the rest of the tree.

let leaf_count = 4
let target = 1
let distances = [
    [0, 13, 21, 22],
    [13, 0, 12, 13],
    [21, 12, 0, 13],
    [22, 13, 13, 0],
]

# For any two other leaves i and k, the paths from j to each of them share the
# limb and then diverge, so (D[i][j] + D[j][k] - D[i][k]) / 2 is the limb plus
# however much further the two paths run together. Taking the minimum over all
# pairs picks the pair that diverges immediately, leaving the limb alone.
#
# This is what makes BA7C possible: knowing the limb length, it can be subtracted
# off and the leaf removed, shrinking the problem by one.
let candidates = range(0, leaf_count)
    |> filter(|i| i != target)
    |> flat_map(|i| range(0, leaf_count)
        |> filter(|k| k != target and k != i)
        |> map(|k| (distances[i][target] + distances[target][k] - distances[i][k]) / 2))

let limb = min(candidates)

println("Result:   " + str(limb))
println("Expected: 2")

fn test_ba7b_limb_length() {
    assert limb == 2, "BA7B: got " + str(limb)
    # The limb can never be longer than half of any distance from j, and never
    # negative in an additive matrix.
    assert limb >= 0, "BA7B: a limb cannot have negative length"
    for i in range(0, leaf_count) {
        if i != target {
            assert limb <= distances[target][i],
                "BA7B: the limb cannot exceed the distance to leaf " + str(i)
        }
    }
    # Leaf 0's limb in the same tree is 11, which BA7A's tree confirms.
    let for_leaf_zero = range(1, leaf_count)
        |> flat_map(|i| range(1, leaf_count)
            |> filter(|k| k != i)
            |> map(|k| (distances[i][0] + distances[0][k] - distances[i][k]) / 2))
    assert min(for_leaf_zero) == 11, "BA7B: leaf 0 hangs off an 11-long limb"
}

# Rosalind: BA7C — Implement AdditivePhylogeny
# https://rosalind.info/problems/ba7c/
#
# Given: An integer n and an additive n x n distance matrix.
# Return: The weighted adjacency list of the simple tree fitting the matrix.

let leaf_count = 4
let distances = [
    [0, 13, 21, 22],
    [13, 0, 12, 13],
    [21, 12, 0, 13],
    [22, 13, 13, 0],
]

# The exact inverse of BA7A: given the distances, recover the tree. It works by
# shrinking the problem — compute a leaf's limb length as in BA7B, subtract it
# from that leaf's row and column, and the leaf now sits at distance zero from
# its attachment point, so it can be removed. Solve the smaller matrix, then hang
# the leaf back on at the right place.
#
# This is exact only because the matrix is additive: some pair of remaining
# leaves must have the attachment point exactly on the path between them, which
# is what says where to graft.
let adjacency = {}
let next_id = leaf_count

fn connect(graph, a, b, weight) {
    let entry = { to_node: b, weight: weight }
    if contains(keys(graph), a) {
        graph[a] = push(graph[a], entry)
    } else {
        graph[a] = [entry]
    }
    graph
}

fn limb_length_of(matrix, leaf, size) {
    range(0, size)
        |> filter(|i| i != leaf)
        |> flat_map(|i| range(0, size)
            |> filter(|k| k != leaf and k != i)
            |> map(|k| (matrix[i][leaf] + matrix[leaf][k] - matrix[i][k]) / 2))
        |> min()
}

# Distance along the tree between two nodes, used to find where to graft.
fn tree_distance(graph, start, finish) {
    let seen = {}
    seen[start] = 0
    let frontier = [start]
    while len(frontier) > 0 {
        let node = frontier[0]
        frontier = slice(frontier, 1, len(frontier))
        if contains(keys(graph), node) {
            for edge in graph[node] {
                if contains(keys(seen), edge.to_node) == false {
                    seen[edge.to_node] = seen[node] + edge.weight
                    frontier = push(frontier, edge.to_node)
                }
            }
        }
    }
    seen[finish]
}

# The path of nodes between two leaves, so the graft point can be located on it.
fn path_between(graph, start, finish) {
    let came_from = {}
    let seen = { }
    seen[start] = true
    let frontier = [start]
    while len(frontier) > 0 {
        let node = frontier[0]
        frontier = slice(frontier, 1, len(frontier))
        if contains(keys(graph), node) {
            for edge in graph[node] {
                if contains(keys(seen), edge.to_node) == false {
                    seen[edge.to_node] = true
                    came_from[edge.to_node] = node
                    frontier = push(frontier, edge.to_node)
                }
            }
        }
    }
    let walk = [finish]
    let at = finish
    while at != start {
        at = came_from[at]
        walk = push(walk, at)
    }
    reverse(walk)
}

fn build(matrix, size, graph, fresh) {
    if size == 2 {
        let g = connect(connect(graph, "0", "1", matrix[0][1]), "1", "0", matrix[0][1])
        return { graph: g, next: fresh }
    }

    let leaf = size - 1
    let limb = limb_length_of(matrix, leaf, size)
    # Trim the limb off, so the leaf sits exactly on its attachment point. Built
    # as a new matrix rather than edited in place — `m[i][j] = x` is not
    # something the language offers, only `m[i] = row`.
    let trimmed = range(0, size) |> map(|r| range(0, size) |> map(|c| {
        let touches_leaf = (r == leaf and c != leaf) or (c == leaf and r != leaf)
        if touches_leaf then matrix[r][c] - limb else matrix[r][c]
    }))

    # Two leaves whose path passes through the attachment point.
    let found = { i: 0, k: 0, at: 0 }
    for i in range(0, size - 1) {
        for k in range(0, size - 1) {
            if i != k and trimmed[i][k] == trimmed[i][leaf] + trimmed[leaf][k] {
                found = { i: i, k: k, at: trimmed[i][leaf] }
            }
        }
    }

    let smaller = build(matrix, size - 1, graph, fresh)
    let g = smaller.graph
    let id = smaller.next

    # Walk from i towards k until the attachment distance is reached; the graft
    # point is either an existing node or a new one splitting an edge.
    let walk = path_between(g, str(found.i), str(found.k))
    let travelled = 0
    let attach = walk[0]
    let previous = walk[0]
    for step in range(1, len(walk)) {
        if travelled < found.at {
            previous = walk[step - 1]
            travelled = travelled + tree_distance(g, walk[step - 1], walk[step])
            attach = walk[step]
        }
    }

    if travelled == found.at {
        # Lands exactly on an existing node.
        g = connect(connect(g, attach, str(leaf), limb), str(leaf), attach, limb)
        return { graph: g, next: id }
    }

    # Otherwise split the edge previous->attach with a new internal node.
    let overshoot = travelled - found.at
    let edge_weight = tree_distance(g, previous, attach)
    let new_node = str(id)
    g[previous] = g[previous] |> filter(|e| e.to_node != attach)
    g[attach] = g[attach] |> filter(|e| e.to_node != previous)
    g = connect(connect(g, previous, new_node, edge_weight - overshoot),
                new_node, previous, edge_weight - overshoot)
    g = connect(connect(g, new_node, attach, overshoot), attach, new_node, overshoot)
    g = connect(connect(g, new_node, str(leaf), limb), str(leaf), new_node, limb)
    { graph: g, next: id + 1 }
}

let built = build(distances, leaf_count, adjacency, leaf_count)
let tree = built.graph

let listed = sort(keys(tree)) |> flat_map(|node|
    tree[node] |> map(|edge| node + "->" + edge.to_node + ":" + str(edge.weight)))

println("Result:")
for line in listed { println("  " + line) }
println("Expected: 0->4:11, 1->4:2, 2->5:6, 3->5:7, 4->5:4")

fn test_ba7c_additive_phylogeny() {
    # The real requirement: the tree reproduces the matrix it was built from.
    # That is checkable directly, and stronger than matching one printed layout,
    # since the internal node numbering is not unique.
    for i in range(0, leaf_count) {
        for j in range(0, leaf_count) {
            assert tree_distance(tree, str(i), str(j)) == distances[i][j],
                "BA7C: leaves " + str(i) + " and " + str(j) + " are "
                    + str(tree_distance(tree, str(i), str(j)))
                    + " apart in the tree but " + str(distances[i][j]) + " in the matrix"
        }
    }
    # A tree with n leaves and only internal nodes of degree 3 has n - 2 of them.
    let internal = keys(tree) |> filter(|node| int(node) >= leaf_count)
    assert len(internal) == leaf_count - 2,
        "BA7C: expected " + str(leaf_count - 2) + " internal nodes, got " + str(len(internal))
}

# Rosalind: BA7D — Implement UPGMA
# https://rosalind.info/problems/ba7d/
#
# Given: An integer n and an n x n distance matrix.
# Return: The adjacency list of the ultrametric tree UPGMA builds.

let leaf_count = 4
let distances = [
    [0, 20, 17, 11],
    [20, 0, 20, 13],
    [17, 20, 0, 10],
    [11, 13, 10, 0],
]

# Hierarchical clustering with heights attached. Each new node is placed at half
# the distance between the two clusters it joins, so every leaf ends up the same
# distance from the root — that is what "ultrametric" means, and it amounts to
# assuming a molecular clock: every lineage evolving at the same rate.
#
# That assumption is often wrong, which is what BA7E's neighbour joining exists to
# avoid. UPGMA is here to show what the simpler assumption costs.
let clusters = range(0, leaf_count) |> map(|i| [i])
let sizes = range(0, leaf_count) |> map(|_| 1)
let ages = range(0, leaf_count) |> map(|_| 0.0)
let live = range(0, leaf_count)
let separation = distances |> map(|line| line |> map(|d| d * 1.0))
let next_id = leaf_count
let adjacency = {}

fn connect(graph, a, b, weight) {
    let entry = { to_node: str(b), weight: round(weight, 3) }
    if contains(keys(graph), str(a)) {
        graph[str(a)] = push(graph[str(a)], entry)
    } else {
        graph[str(a)] = [entry]
    }
    graph
}

while len(live) > 1 {
    # The closest pair of clusters, by average distance.
    let best_a = live[0]
    let best_b = live[1]
    let best = separation[best_a][best_b]
    for a in live {
        for b in live {
            if a < b and separation[a][b] < best {
                best = separation[a][b]
                best_a = a
                best_b = b
            }
        }
    }

    # The new node sits half the distance up, and each child's limb is whatever
    # is left after its own age.
    let age = best / 2
    adjacency = connect(adjacency, next_id, best_a, age - ages[best_a])
    adjacency = connect(adjacency, best_a, next_id, age - ages[best_a])
    adjacency = connect(adjacency, next_id, best_b, age - ages[best_b])
    adjacency = connect(adjacency, best_b, next_id, age - ages[best_b])

    # Average distance from the merged cluster to everything else, weighted by
    # how many leaves each side holds.
    let merged_size = sizes[best_a] + sizes[best_b]
    let row = range(0, next_id + 1) |> map(|_| 0.0)
    for other in live {
        if other != best_a and other != best_b {
            row[other] = (separation[best_a][other] * sizes[best_a]
                        + separation[best_b][other] * sizes[best_b]) / merged_size
        }
    }
    separation = push(separation, row)
    for other in live {
        if other != best_a and other != best_b {
            separation[other] = push(separation[other], row[other])
        }
    }

    sizes = push(sizes, merged_size)
    ages = push(ages, age)
    live = (live |> filter(|node| node != best_a and node != best_b)) + [next_id]
    next_id = next_id + 1
}

let listed = sort(keys(adjacency)) |> flat_map(|node|
    adjacency[node] |> map(|edge| node + "->" + edge.to_node + ":" + str(round(edge.weight, 3))))

println("Result:")
for line in listed { println("  " + line) }
println("Expected: 0->5:7.000, 1->6:8.833, 2->4:5.000, 3->4:5.000, 4->5:2.000, 5->6:1.833")

fn test_ba7d_upgma() {
    # Every edge, checked by name so ordering does not matter.
    # Rosalind prints these to three decimals; BioLang prints the shortest exact
    # form, so 7.000 shows as 7.0. The values are the same.
    let want = ["0->5:7.0", "5->0:7.0", "1->6:8.833", "6->1:8.833", "2->4:5.0", "4->2:5.0",
                "3->4:5.0", "4->3:5.0", "4->5:2.0", "5->4:2.0", "5->6:1.833", "6->5:1.833"]
    assert len(listed) == len(want),
        "BA7D: expected " + str(len(want)) + " edges, got " + str(len(listed))
    for edge in want {
        assert contains(listed, edge), "BA7D: missing edge " + edge
    }
    # Ultrametric is the claim worth checking: every leaf sits the same distance
    # from the root. That is the molecular-clock assumption made concrete, and
    # what BA7E drops.
    let root_age = ages[len(ages) - 1]
    assert abs(root_age - 8.833) < 5e-4, "BA7D: the root sits at " + str(round(root_age, 3))
    assert abs((7.0 + 1.833) - root_age) < 5e-3, "BA7D: leaf 0 reaches it via node 5"
    assert abs((5.0 + 2.0 + 1.833) - root_age) < 5e-3, "BA7D: leaf 2 via nodes 4 and 5"
    assert abs(8.833 - root_age) < 5e-4, "BA7D: leaf 1 reaches it directly"
}

# Rosalind: BA7E — Implement the Neighbor Joining Algorithm
# https://rosalind.info/problems/ba7e/
#
# Given: An integer n and an n x n distance matrix.
# Return: The adjacency list of the tree neighbour joining builds.

let leaf_count = 4
let distances = [
    [0, 23, 27, 20],
    [23, 0, 30, 28],
    [27, 30, 0, 30],
    [20, 28, 30, 0],
]

# UPGMA joins whichever pair is closest, which is wrong whenever two lineages
# evolve at different rates — a fast-evolving leaf looks distant from its true
# relatives and gets attached elsewhere. Neighbour joining corrects each distance
# by how far each leaf is from *everything* else before comparing, so a
# uniformly-distant leaf is no longer penalised. It makes no molecular-clock
# assumption, and the tree it returns is unrooted.
fn neighbour_matrix(working, live) {
    let n = len(live)
    let totals = {}
    for i in live { totals[str(i)] = live |> map(|j| working[i][j]) |> sum() }
    let corrected = {}
    for i in live {
        for j in live {
            if i != j {
                corrected[str(i) + "," + str(j)] =
                    (n - 2) * working[i][j] - totals[str(i)] - totals[str(j)]
            }
        }
    }
    { corrected: corrected, totals: totals }
}

let working = distances |> map(|line| line |> map(|d| d * 1.0))
let live = range(0, leaf_count)
let next_id = leaf_count
let adjacency = {}

fn connect(graph, a, b, weight) {
    let entry = { to_node: str(b), weight: round(weight, 3) }
    if contains(keys(graph), str(a)) {
        graph[str(a)] = push(graph[str(a)], entry)
    } else {
        graph[str(a)] = [entry]
    }
    graph
}

while len(live) > 2 {
    let built = neighbour_matrix(working, live)
    let corrected = built.corrected
    let totals = built.totals

    let best_a = live[0]
    let best_b = live[1]
    let best = corrected[str(best_a) + "," + str(best_b)]
    for a in live {
        for b in live {
            if a != b and corrected[str(a) + "," + str(b)] < best {
                best = corrected[str(a) + "," + str(b)]
                best_a = a
                best_b = b
            }
        }
    }

    let n = len(live)
    # How lopsided the pair is — this is what lets the two limbs differ, which
    # UPGMA cannot express.
    let delta = (totals[str(best_a)] - totals[str(best_b)]) / (n - 2)
    let limb_a = (working[best_a][best_b] + delta) / 2
    let limb_b = (working[best_a][best_b] - delta) / 2

    let row = range(0, next_id + 1) |> map(|_| 0.0)
    for other in live {
        if other != best_a and other != best_b {
            let joined = working[best_a][other] + working[best_b][other]
            row[other] = (joined - working[best_a][best_b]) / 2
        }
    }
    working = push(working, row)
    for other in live {
        if other != best_a and other != best_b {
            working[other] = push(working[other], row[other])
        }
    }

    adjacency = connect(adjacency, next_id, best_a, limb_a)
    adjacency = connect(adjacency, best_a, next_id, limb_a)
    adjacency = connect(adjacency, next_id, best_b, limb_b)
    adjacency = connect(adjacency, best_b, next_id, limb_b)

    live = (live |> filter(|node| node != best_a and node != best_b)) + [next_id]
    next_id = next_id + 1
}

# Two left: join them with the distance between them.
let last_a = live[0]
let last_b = live[1]
adjacency = connect(adjacency, last_a, last_b, working[last_a][last_b])
adjacency = connect(adjacency, last_b, last_a, working[last_a][last_b])

let listed = sort(keys(adjacency)) |> flat_map(|node|
    adjacency[node] |> map(|edge| node + "->" + edge.to_node + ":" + str(round(edge.weight, 3))))

println("Result:")
for line in listed { println("  " + line) }
println("Expected: 0->4:8, 1->5:13.5, 2->5:16.5, 3->4:12, 4->5:2")

fn test_ba7e_neighbour_joining() {
    let want = ["0->4:8.0", "4->0:8.0", "1->5:13.5", "5->1:13.5", "2->5:16.5", "5->2:16.5",
                "3->4:12.0", "4->3:12.0", "4->5:2.0", "5->4:2.0"]
    assert len(listed) == len(want),
        "BA7E: expected " + str(len(want)) + " edges, got " + str(len(listed))
    for edge in want { assert contains(listed, edge), "BA7E: missing edge " + edge }

    # This matrix is *not* additive, which the four-point condition shows: for an
    # additive matrix the two largest of the three pairings must be equal, and
    # here they are 55, 53 and 50. So no tree reproduces it exactly, and neighbour
    # joining returns a best fit rather than an exact answer — worth asserting,
    # because the natural expectation is that the distances come back unchanged.
    let pairing_one = distances[0][1] + distances[2][3]
    let pairing_two = distances[0][2] + distances[1][3]
    let pairing_three = distances[0][3] + distances[1][2]
    assert pairing_one == 53 and pairing_two == 55 and pairing_three == 50,
        "BA7E: the three pairings are 53, 55, 50"
    let two_largest_agree = pairing_one == pairing_two or pairing_one == pairing_three
        or pairing_two == pairing_three
    assert two_largest_agree == false, "BA7E: so the working is not additive"

    # Some distances the tree does get exactly right, and none is off by much.
    assert (8.0 + 12.0) == distances[0][3], "BA7E: 0 to 3 is 8 + 12 = 20, exactly"
    assert (13.5 + 16.5) == distances[1][2], "BA7E: 1 to 2 is 30, exactly"
    assert abs((8.0 + 2.0 + 13.5) - distances[0][1]) <= 0.5,
        "BA7E: 0 to 1 comes out at 23.5 against 23 — the cost of non-additivity"
    assert abs((8.0 + 2.0 + 16.5) - distances[0][2]) <= 0.5, "BA7E: 0 to 2 within half"

    # Limbs differ within a pair, which is exactly what UPGMA cannot express.
    assert 13.5 != 16.5, "BA7E: the two limbs off node 5 have different lengths"
}

# Rosalind: BA7F — Implement SmallParsimony
# https://rosalind.info/problems/ba7f/
#
# Given: An integer n and a rooted binary tree with n leaves labelled by DNA
# strings.
# Return: The minimum parsimony score, and a labelling of the internal nodes
# achieving it.

let leaf_count = 4
let leaves = {
    "0": "CAAATCCC",
    "1": "ATTGCGAC",
    "2": "CTGCGCTG",
    "3": "ATGGACGA",
}
let children = {
    "4": ["0", "1"],
    "5": ["2", "3"],
    "6": ["4", "5"],
}
let root = "6"

# Sankoff's algorithm. Every column of the alignment is independent — a mutation
# at one site says nothing about another — so each is solved separately and the
# scores added. Within a column, the cheapest cost of a subtree given its root's
# base is the sum over children of (their cheapest cost, plus one if the base has
# to change). Working leaves-upwards means each child is already solved.
#
# Choosing greedily from the leaves would not work: a base that looks locally
# cheap at one node can force two changes higher up.
let bases = ["A", "C", "G", "T"]
let width = len(leaves["0"])

fn is_leaf(node, leaf_map) { contains(keys(leaf_map), node) }

# Cheapest cost of the subtree at `node` for each possible base at `node`.
fn score_column(node, position, leaf_map, child_map, alphabet) {
    if is_leaf(node, leaf_map) {
        let here = substr(leaf_map[node], position, 1)
        # Same record shape as an internal node, with nothing below it.
        return {
            costs: alphabet |> map(|b| if b == here then 0 else 1000000),
            picks: {}, solved: [], kids: [],
        }
    }
    let kids = child_map[node]
    let solved = kids |> map(|kid| score_column(kid, position, leaf_map, child_map, alphabet))
    let picks = {}
    let costs = range(0, len(alphabet)) |> map(|mine| {
        range(0, len(kids)) |> map(|k| {
            # The child's own cost plus one if its best base differs from mine.
            let options = range(0, len(alphabet))
                |> map(|theirs| solved[k].costs[theirs] + (if theirs == mine then 0 else 1))
            let chosen = argmin(options)
            picks[str(mine) + "," + str(k)] = chosen
            options[chosen]
        }) |> sum()
    })
    { costs: costs, picks: picks, solved: solved, kids: kids }
}

# Walk back down, fixing each node's base from its parent's choice. Leaves are
# already labelled by the problem, so only internal nodes are written.
fn assign(node, tree, chosen_index, labels, alphabet) {
    labels[node] = labels[node] + alphabet[chosen_index]
    for k in range(0, len(tree.kids)) {
        if len(tree.solved[k].kids) > 0 {
            labels = assign(tree.kids[k], tree.solved[k],
                            tree.picks[str(chosen_index) + "," + str(k)], labels, alphabet)
        }
    }
    labels
}

let labels = {}
for node in keys(children) { labels[node] = "" }
let total = 0

for position in range(0, width) {
    let solved = score_column(root, position, leaves, children, bases)
    let best = argmin(solved.costs)
    total = total + solved.costs[best]
    labels = assign(root, solved, best, labels, bases)
}

# Every node's string: leaves as given, internals as just computed.
let named = {}
for node in keys(leaves) { named[node] = leaves[node] }
for node in keys(children) { named[node] = labels[node] }

fn hamming(a, b) { range(0, len(a)) |> count_if(|i| substr(a, i, 1) != substr(b, i, 1)) }

let listed_edges = sort(keys(children)) |> flat_map(|parent|
    children[parent] |> flat_map(|kid| [
        named[parent] + "->" + named[kid] + ":" + str(hamming(named[parent], named[kid])),
        named[kid] + "->" + named[parent] + ":" + str(hamming(named[parent], named[kid])),
    ]))

println("Result:   " + str(total))
for line in listed_edges { println("  " + line) }
println("Expected: 16")

fn test_ba7f_small_parsimony() {
    assert total == 16, "BA7F: scored " + str(total)
    # The score has to equal the changes actually shown on the tree — a score
    # without a matching labelling is the usual bug here.
    let shown = sort(keys(children)) |> flat_map(|parent|
        children[parent] |> map(|kid| hamming(named[parent], named[kid]))) |> sum()
    assert shown == total,
        "BA7F: the labelling shows " + str(shown) + " changes but the score is " + str(total)
    # Every internal label is a real DNA string of the right length.
    for node in keys(children) {
        assert len(named[node]) == width, "BA7F: " + node + " has the wrong length"
        assert (chars(named[node]) |> count_if(|b| contains(bases, b) == false)) == 0,
            "BA7F: " + named[node] + " contains a non-base"
    }
    # And no labelling can do better, which the published answer confirms.
    assert total <= 16, "BA7F: 16 is the published minimum"
}

# Rosalind: BA7G — Adapt SmallParsimony to Unrooted Trees
# https://rosalind.info/problems/ba7g/
#
# Given: An unrooted binary tree whose leaves are labelled by DNA strings.
# Return: The minimum parsimony score and a labelling achieving it.

let leaves = {
    "0": "TCGGCCAA",
    "1": "CCTGGCTG",
    "2": "CACAGGAT",
    "3": "TGAGTACC",
}
# The unrooted tree: 0 and 1 hang off node 4, 2 and 3 off node 5, and 4-5 join
# the two halves.
let unrooted_edge = ["4", "5"]

# Parsimony counts changes along edges, and an edge's cost does not depend on
# which direction it is read — so the score of an unrooted tree is the same
# whichever edge a root is placed on. That is the whole adaptation: hang a
# temporary root in the middle of any edge, run BA7F unchanged, then delete the
# root and reconnect the two nodes it separated.
#
# The labelling can differ depending on where the root goes; the score cannot.
let children = {
    "4": ["0", "1"],
    "5": ["2", "3"],
    "6": ["4", "5"],
}
let root = "6"
let bases = ["A", "C", "G", "T"]
let width = len(leaves["0"])

fn is_leaf(node, leaf_map) { contains(keys(leaf_map), node) }

fn score_column(node, position, leaf_map, child_map, alphabet) {
    if is_leaf(node, leaf_map) {
        let here = substr(leaf_map[node], position, 1)
        return {
            costs: alphabet |> map(|b| if b == here then 0 else 1000000),
            picks: {}, solved: [], kids: [],
        }
    }
    let kids = child_map[node]
    let solved = kids |> map(|kid| score_column(kid, position, leaf_map, child_map, alphabet))
    let picks = {}
    let costs = range(0, len(alphabet)) |> map(|mine| {
        range(0, len(kids)) |> map(|k| {
            let options = range(0, len(alphabet))
                |> map(|theirs| solved[k].costs[theirs] + (if theirs == mine then 0 else 1))
            let chosen = argmin(options)
            picks[str(mine) + "," + str(k)] = chosen
            options[chosen]
        }) |> sum()
    })
    { costs: costs, picks: picks, solved: solved, kids: kids }
}

fn assign(node, tree, chosen_index, labels, alphabet) {
    labels[node] = labels[node] + alphabet[chosen_index]
    for k in range(0, len(tree.kids)) {
        if len(tree.solved[k].kids) > 0 {
            labels = assign(tree.kids[k], tree.solved[k],
                            tree.picks[str(chosen_index) + "," + str(k)], labels, alphabet)
        }
    }
    labels
}

let labels = {}
for node in keys(children) { labels[node] = "" }
let rooted_total = 0
for position in range(0, width) {
    let solved = score_column(root, position, leaves, children, bases)
    let best = argmin(solved.costs)
    rooted_total = rooted_total + solved.costs[best]
    labels = assign(root, solved, best, labels, bases)
}

let named = {}
for node in keys(leaves) { named[node] = leaves[node] }
for node in keys(children) { named[node] = labels[node] }

fn hamming(a, b) { range(0, len(a)) |> count_if(|i| substr(a, i, 1) != substr(b, i, 1)) }

# Remove the temporary root: its two children are joined directly, and the two
# edges to the root become one.
let final_edges = [
    { a: "0", b: "4" }, { a: "1", b: "4" },
    { a: "2", b: "5" }, { a: "3", b: "5" },
    { a: unrooted_edge[0], b: unrooted_edge[1] },
]

let total = final_edges |> map(|e| hamming(named[e.a], named[e.b])) |> sum()

let listed = final_edges |> flat_map(|e| [
    named[e.a] + "->" + named[e.b] + ":" + str(hamming(named[e.a], named[e.b])),
    named[e.b] + "->" + named[e.a] + ":" + str(hamming(named[e.a], named[e.b])),
])

println("Result:   " + str(total))
for line in listed { println("  " + line) }
println("Expected: 17")

fn test_ba7g_unrooted_small_parsimony() {
    assert total == 17, "BA7G: scored " + str(total)
    # The score shown on the edges must be the score reported.
    let shown = final_edges |> map(|e| hamming(named[e.a], named[e.b])) |> sum()
    assert shown == total, "BA7G: the labelling shows " + str(shown) + " changes"
    # Rooting adds no cost of its own: the rooted run scores the same as the
    # unrooted tree it stands for. That equality is the whole justification for
    # solving the problem this way.
    assert rooted_total == total,
        "BA7G: rooted scored " + str(rooted_total) + " but unrooted scored " + str(total)
    # Four leaves in an unrooted binary tree means two internal nodes and five
    # edges.
    assert len(final_edges) == 5, "BA7G: an unrooted tree with 4 leaves has 5 edges"
    for node in ["4", "5"] {
        assert len(named[node]) == width, "BA7G: " + node + " has the wrong length"
    }
}