二分ヒープによるハフマン符号の効率化

二分ヒープを用いることでハフマン符号を効率よく計算できるようにしました

なぜ二分ヒープを用いるのか

ハフマン木を作るとき、出現回数が最小要素と二番目に小さい要素を求める必要があります。この操作は線形探索ではO(n)かかりますが、二分ヒープを用いることでO(log n)で求めることができます。

二分ヒープの実装

wikiまんまですがこんな感じで実装しました。のちにハフマン木を構築することを考えて若干柔軟な実装にしてあります。

class Heap {
    constructor(comp = (x, y) => x < y) {
        this.node = [];
        this.compare = comp; // 応用性を持たせるための比較関数。これで最小ヒープ、最大ヒープをこのクラスのみで構築可能
        // compは左辺 < 右辺ならばtrueを返すという関数。最小ヒープにしたいならば左辺 > 右辺とすればよい
    }
    push(...values) {
        for (const value of values) {
            this.node.push(value);
            let x = this.node.length - 1, parent = 0 | (x - 1) / 2; // 現在のノードと親ノード
            while (x > 0 && this.compare(this.node[parent], this.node[x])) {
                // 親ノードと現在のノードを入れ替える
                const tmp = this.node[x];
                this.node[x] = this.node[parent];
                this.node[parent] = tmp;

                x = parent;
                parent = 0 | (x - 1) / 2;
            }
        }
    }
    pop() {
        const res = this.node[0];
        this.node[0] = this.node[this.node.length - 1];　// 末尾要素を根に持ってくる
        this.node.splice(this.node.length - 1, 1); // 末尾要素の削除
        let x = 0, left = x * 2 + 1, right = x * 2 + 2; // 右の子
        const length = this.node.length, comp = this.compare;
        while (left < length && comp(this.node[x], this.node[left]) || right < length && comp(this.node[x], this.node[right])) {
            const max = right < length && comp(this.node[left], this.node[right]) ? right : left; // 右の子がないことがあるため、このようにした
            // 比較関数により、より"大きい"とされた方と入れ替える
            const tmp = this.node[x];
            this.node[x] = this.node[max];
            this.node[max] = tmp;

            x = max;
            left = x * 2 + 1;
            right = x * 2 + 2;
        }
        return res;
    }
}

ハフマン木を構築する

上のHeapクラスを用いてハフマン木を構築します。

class HuffmanNode {
    constructor(symbol, weight) {
        this.symbol = symbol;
        this.weight = weight;
        this.child = [];
    }
    addLeft(node) {this.child[0] = node;}
    addRight(node) {this.child[1] = node;}
}
function count(str) {
    let counter = {};
    for (const s of str) {
        if (!counter[s]) counter[s] = 0;
        counter[s]++;
    }
    const res = [];
    for (const key of Object.keys(counter)) {
        res.push([key, counter[key]]);
    }
    return res;
}
function createHuffmanTree(pairs) {
    // pairs:= [['a', 2], ['b', 2], ...]
    pairs = pairs.map(item => new HuffmanNode(item[0], item[1]));
    const heap = new Heap((a, b) => a.weight > b.weight);
    heap.push(...pairs);
    while (heap.node.length > 1) {
        const a = heap.pop();
        const b = heap.pop();
        const parent = new HuffmanNode(null, a.weight + b.weight);
        parent.addLeft(a);
        parent.addRight(b);
        heap.push(parent);
    }
    return heap.pop();
}
createHuffmanTree(count('Lorem ipsum dolor sit amet, consectetur adipiscing elit'));

前回の記事から若干Nodeクラスにも変更を加えてあります。前回はNodeにはsymbolのみを持たせていましたが、今回はweightもNodeに持たせました。

ハフマン符号を求める

Nodeに変更を加えたので、前回の記事の関数が使えなくなりました。そこで若干の書き直しを行いました。

function HuffmanCoding(tree, dic = {}, s = '') {
    if (tree.symbol !== null) dic[tree.symbol] = s;
    if (tree.child[0]) dic = HuffmanCoding(tree.child[0], dic, s + '0');
    if (tree.child[1]) dic = HuffmanCoding(tree.child[1], dic, s + '1');
    return dic;
}

.value を .symbolに直しただけです。

全体のコード

class Heap {
    constructor(comp = (x, y) => x < y) {
        this.node = [];
        this.compare = comp;
    }
    push(...values) {
        for (const value of values) {
            this.node.push(value);
            let x = this.node.length - 1, parent = 0 | (x - 1) / 2;
            while (x > 0 && !this.compare(this.node[x], this.node[parent])) {
                const tmp = this.node[x];
                this.node[x] = this.node[parent];
                this.node[parent] = tmp;

                x = parent;
                parent = 0 | (x - 1) / 2;
            }
        }
    }
    pop() {
        const res = this.node[0];
        this.node[0] = this.node[this.node.length - 1];　// 末尾要素を根に持ってくる
        this.node.splice(this.node.length - 1, 1); // 末尾要素の削除
        let x = 0, left = x * 2 + 1, right = x * 2 + 2; // 右の子
        const length = this.node.length, comp = this.compare;
        while (left < length && comp(this.node[x], this.node[left]) || right < length && comp(this.node[x], this.node[right])) {
            const max = right < length && comp(this.node[left], this.node[right]) ? right : left;
            const tmp = this.node[x];
            this.node[x] = this.node[max];
            this.node[max] = tmp;

            x = max;
            left = x * 2 + 1;
            right = x * 2 + 2;
        }
        return res;
    }
}
class HuffmanNode {
    constructor(symbol, weight) {
        this.symbol = symbol;
        this.weight = weight;
        this.child = [];
    }
    addLeft(node) {this.child[0] = node;}
    addRight(node) {this.child[1] = node;}
}
function count(str) {
    let counter = {};
    for (const s of str) {
        if (!counter[s]) counter[s] = 0;
        counter[s]++;
    }
    const res = [];
    for (const key of Object.keys(counter)) {
        res.push([key, counter[key]]);
    }
    return res;
}
function createHuffmanTree(pairs) {
    // pairs:= [['a', 2], ['b', 2], ...]
    pairs = pairs.map(item => new HuffmanNode(item[0], item[1]));
    const heap = new Heap((a, b) => a.weight > b.weight);
    heap.push(...pairs);
    while (heap.node.length > 1) {
        const a = heap.pop();
        const b = heap.pop();
        const parent = new HuffmanNode(null, a.weight + b.weight);
        parent.addLeft(a);
        parent.addRight(b);
        heap.push(parent);
    }
    return heap.pop();
}
function HuffmanCoding(tree, dic = {}, s = '') {
    if (tree.symbol !== null) dic[tree.symbol] = s;
    if (tree.child[0]) dic = HuffmanCoding(tree.child[0], dic, s + '0');
    if (tree.child[1]) dic = HuffmanCoding(tree.child[1], dic, s + '1');
    return dic;
}
const test = 'Lorem ipsum dolor sit amet, consectetur adipiscing elit';
console.log(HuffmanCoding(createHuffmanTree(count(test))));
*