Huffman Coding

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
import heapq


@dataclass
class Node:
    frequency: int
    symbol: int | None = None
    left: "Node | None" = None
    right: "Node | None" = None

    @property
    def is_leaf(self) -> bool:
        return self.left is None and self.right is None


def build_tree(data: bytes) -> Node | None:
    """Build a deterministic tree over byte symbols 0..255."""
    heap = []
    serial = 0
    for symbol, frequency in sorted(Counter(data).items()):
        node = Node(frequency, symbol)
        heapq.heappush(heap, (frequency, symbol, serial, node))
        serial += 1

    # Invariant: heap entries are roots of disjoint subtrees whose leaves
    # partition the source alphabet; their frequencies sum to len(data).
    while len(heap) > 1:
        left_item = heapq.heappop(heap)
        right_item = heapq.heappop(heap)
        left, right = left_item[3], right_item[3]
        parent = Node(left.frequency + right.frequency, None, left, right)
        minimum_symbol = min(left_item[1], right_item[1])
        heapq.heappush(
            heap, (parent.frequency, minimum_symbol, serial, parent)
        )
        serial += 1
    return heap[0][3] if heap else None


def build_codes(root: Node | None) -> dict[int, str]:
    codes: dict[int, str] = {}

    def visit(node: Node, prefix: str) -> None:
        if node.is_leaf:
            assert node.symbol is not None
            codes[node.symbol] = prefix or "0"
            return
        assert node.left is not None and node.right is not None
        visit(node.left, prefix + "0")
        visit(node.right, prefix + "1")

    if root is not None:
        visit(root, "")
    return codes


def encode(data: bytes) -> tuple[str, dict[int, str], Node | None]:
    root = build_tree(data)
    codes = build_codes(root)
    return "".join(codes[symbol] for symbol in data), codes, root


def decode(bits: str, root: Node | None, expected_symbols: int) -> bytes:
    """Reject invalid bits, truncated paths, extras, and bad singleton data."""
    if expected_symbols < 0:
        raise ValueError("negative symbol count")
    if root is None:
        if bits or expected_symbols:
            raise ValueError("nonempty payload without a tree")
        return b""
    if root.is_leaf:
        if bits != "0" * expected_symbols:
            raise ValueError("invalid single-symbol payload")
        assert root.symbol is not None
        return bytes([root.symbol]) * expected_symbols

    output = bytearray()
    current = root
    for bit in bits:
        if bit not in "01":
            raise ValueError("payload contains a non-bit")
        current = current.left if bit == "0" else current.right
        if current is None:
            raise ValueError("bit path is absent from tree")
        if current.is_leaf:
            assert current.symbol is not None
            output.append(current.symbol)
            if len(output) > expected_symbols:
                raise ValueError("payload contains too many symbols")
            current = root
    if current is not root or len(output) != expected_symbols:
        raise ValueError("truncated or incorrectly sized payload")
    return bytes(output)


if __name__ == "__main__":
    for sample in (b"", b"AAAA", b"AABBBCCCC", bytes(range(32))):
        bits, codes, tree = encode(sample)
        assert decode(bits, tree, len(sample)) == sample
    bits, codes, tree = encode(b"AABBBCCCC")
    assert len(bits) == 14
    print(bits, {chr(symbol): code for symbol, code in codes.items()})

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

public class HuffmanCoding {
    static final class Node {
        final long frequency;
        final int symbol; // -1 for internal nodes
        final Node left;
        final Node right;

        Node(long frequency, int symbol, Node left, Node right) {
            this.frequency = frequency;
            this.symbol = symbol;
            this.left = left;
            this.right = right;
        }

        boolean isLeaf() {
            return left == null && right == null;
        }
    }

    static final class Entry {
        final long frequency;
        final int minimumSymbol;
        final long serial;
        final Node node;

        Entry(long frequency, int minimumSymbol, long serial, Node node) {
            this.frequency = frequency;
            this.minimumSymbol = minimumSymbol;
            this.serial = serial;
            this.node = node;
        }
    }

    static final class Encoded {
        final String bits;
        final String[] codes;
        final Node root;
        final int originalLength;

        Encoded(String bits, String[] codes, Node root, int originalLength) {
            this.bits = bits;
            this.codes = codes;
            this.root = root;
            this.originalLength = originalLength;
        }
    }

    static Node buildTree(byte[] data) {
        long[] frequencies = new long[256];
        for (byte value : data) frequencies[value & 0xff]++;

        Comparator<Entry> order = Comparator
                .comparingLong((Entry entry) -> entry.frequency)
                .thenComparingInt(entry -> entry.minimumSymbol)
                .thenComparingLong(entry -> entry.serial);
        PriorityQueue<Entry> heap = new PriorityQueue<>(order);
        long serial = 0L;
        for (int symbol = 0; symbol < frequencies.length; symbol++) {
            if (frequencies[symbol] > 0L) {
                Node leaf = new Node(frequencies[symbol], symbol, null, null);
                heap.add(new Entry(frequencies[symbol], symbol, serial++, leaf));
            }
        }

        // Invariant: entries are disjoint subtree roots whose leaf
        // frequencies sum to data.length.
        while (heap.size() > 1) {
            Entry left = heap.remove();
            Entry right = heap.remove();
            long frequency = Math.addExact(left.frequency, right.frequency);
            Node parent = new Node(frequency, -1, left.node, right.node);
            heap.add(new Entry(frequency,
                    Math.min(left.minimumSymbol, right.minimumSymbol),
                    serial++, parent));
        }
        return heap.isEmpty() ? null : heap.remove().node;
    }

    static String[] buildCodes(Node root) {
        String[] codes = new String[256];
        if (root != null) generateCodes(root, "", codes);
        return codes;
    }

    private static void generateCodes(Node node, String prefix, String[] codes) {
        if (node.isLeaf()) {
            codes[node.symbol] = prefix.isEmpty() ? "0" : prefix;
            return;
        }
        generateCodes(node.left, prefix + "0", codes);
        generateCodes(node.right, prefix + "1", codes);
    }

    static Encoded encode(byte[] data) {
        Node root = buildTree(data);
        String[] codes = buildCodes(root);
        StringBuilder bits = new StringBuilder();
        for (byte value : data) bits.append(codes[value & 0xff]);
        return new Encoded(bits.toString(), codes, root, data.length);
    }

    static byte[] decode(String bits, Node root, int expectedSymbols) {
        if (expectedSymbols < 0) throw new IllegalArgumentException("negative size");
        if (root == null) {
            if (!bits.isEmpty() || expectedSymbols != 0) {
                throw new IllegalArgumentException("payload without tree");
            }
            return new byte[0];
        }
        if (root.isLeaf()) {
            if (bits.length() != expectedSymbols
                    || bits.chars().anyMatch(bit -> bit != '0')) {
                throw new IllegalArgumentException("invalid singleton payload");
            }
            byte[] output = new byte[expectedSymbols];
            Arrays.fill(output, (byte) root.symbol);
            return output;
        }

        ByteArrayOutputStream output = new ByteArrayOutputStream();
        Node current = root;
        for (int index = 0; index < bits.length(); index++) {
            char bit = bits.charAt(index);
            if (bit != '0' && bit != '1') {
                throw new IllegalArgumentException("non-bit in payload");
            }
            current = bit == '0' ? current.left : current.right;
            if (current == null) {
                throw new IllegalArgumentException("absent tree path");
            }
            if (current.isLeaf()) {
                output.write(current.symbol);
                if (output.size() > expectedSymbols) {
                    throw new IllegalArgumentException("too many symbols");
                }
                current = root;
            }
        }
        if (current != root || output.size() != expectedSymbols) {
            throw new IllegalArgumentException("truncated or mis-sized payload");
        }
        return output.toByteArray();
    }

    public static void main(String[] args) {
        byte[][] samples = {
            new byte[0],
            "AAAA".getBytes(StandardCharsets.US_ASCII),
            "AABBBCCCC".getBytes(StandardCharsets.US_ASCII)
        };
        for (byte[] sample : samples) {
            Encoded encoded = encode(sample);
            assert Arrays.equals(
                    decode(encoded.bits, encoded.root, encoded.originalLength),
                    sample);
        }
        Encoded encoded = encode(samples[2]);
        assert encoded.bits.length() == 14;
        System.out.println(encoded.bits);
    }
}