neighbor-joining-phylo

Answer

The implementation below performs the Neighbor-Joining algorithm with these steps: 1. Compute row sums R(i) = Σᵢₖ d(i,k) for all active taxa 2. Compute the Q-matrix: Q(i,j) = (n-2)·d(i,j) − R(i) − R(j) 3. Find the pair (i,j) with minimum Q — these are the "neighbors" 4. Compute branch lengths: d(i,u) = ½d(i,j) + ½(R(i)−R(j))/(n−2); d(j,u) = d(i,j) − d(i,u) 5. Update distances: d(u,k) = ½(d(i,k) + d(j,k) − d(i,j)) 6. Replace i,j with u and repeat until 2 nodes remain 7. Connect the last two nodes, each receiving half the distance 8. Output the rooted tree in Newick format

#!/usr/bin/env python3
"""
Neighbor-Joining Phylogenetic Tree Construction

Given a square distance matrix for n taxa, construct a phylogenetic tree
using the Neighbor-Joining (NJ) algorithm. Returns the tree in Newick format.
"""

from __future__ import annotations
from typing import Dict, List, Optional


class Node:
    """A node in the phylogenetic tree."""

    def __init__(self, name: Optional[str] = None, branch_length: float = 0.0):
        self.name = name
        self.branch_length = branch_length
        self.children: List[Node] = []

    def is_leaf(self) -> bool:
        return self.name is not None

    def to_newick(self) -> str:
        """Recursively build the Newick representation."""
        if self.is_leaf():
            return f"{self.name}:{self.branch_length:.10f}"
        children_str = ",".join(c.to_newick() for c in self.children)
        return f"({children_str}):{self.branch_length:.10f}"


def neighbor_joining(names: List[str], distance_matrix: List[List[float]]) -> str:
    """
    Construct a phylogenetic tree using the Neighbor-Joining algorithm.

    Args:
        names: List of n taxa names.
        distance_matrix: n x n symmetric distance matrix (zero diagonal).

    Returns:
        Newick string of the inferred tree with branch lengths.
    """
    n = len(names)
    dist = [row[:] for row in distance_matrix]
    active: List[int] = list(range(n))
    nodes: Dict[int, Node] = {i: Node(name=names[i]) for i in range(n)}
    next_id = n

    while len(active) > 2:
        m = len(active)

        # Step 1: Row sums R(i)
        R = [0.0] * m
        for i in range(m):
            s = 0.0
            ii = active[i]
            for k in range(m):
                if i != k:
                    s += dist[ii][active[k]]
            R[i] = s

        # Step 2: Find minimum Q(i,j)
        min_q = float("inf")
        min_i = min_j = -1
        for i in range(m):
            for j in range(i + 1, m):
                q_val = (m - 2) * dist[active[i]][active[j]] - R[i] - R[j]
                if q_val < min_q:
                    min_q = q_val
                    min_i, min_j = i, j

        # Step 3: Merge closest pair
        i_glob, j_glob = active[min_i], active[min_j]
        d_ij = dist[i_glob][j_glob]
        len_i = 0.5 * d_ij + 0.5 * (R[min_i] - R[min_j]) / (m - 2)
        len_j = d_ij - len_i

        u = Node()
        u.children = [nodes[i_glob], nodes[j_glob]]
        nodes[i_glob].branch_length = len_i
        nodes[j_glob].branch_length = len_j
        u_id = next_id
        next_id += 1
        nodes[u_id] = u

        # Step 4: Update distance matrix
        u_distances: Dict[int, float] = {}
        for k_idx in range(m):
            if k_idx in (min_i, min_j):
                continue
            k_glob = active[k_idx]
            d_uk = 0.5 * (dist[i_glob][k_glob] + dist[j_glob][k_glob] - d_ij)
            u_distances[k_glob] = d_uk

        # Update active list
        i_pos, j_pos = active.index(i_glob), active.index(j_glob)
        active.pop(max(i_pos, j_pos))
        active.pop(min(i_pos, j_pos))
        active.append(u_id)

        # Extend dist matrix for new node
        for row_idx in range(len(dist)):
            dist[row_idx].append(u_distances.get(row_idx, 0.0))
        new_row = [0.0] * (len(dist) + 1)
        for k_glob, d_uk in u_distances.items():
            new_row[k_glob] = d_uk
        dist.append(new_row)

    # Final step: connect last two nodes
    a, b = active[0], active[1]
    nodes[a].branch_length = dist[a][b] / 2.0
    nodes[b].branch_length = dist[a][b] / 2.0

    root = Node()
    root.children = [nodes[a], nodes[b]]
    return root.to_newick() + ";"