题解 | 哈夫曼树

#include <iostream>
#include <set>

using std::cout;
using std::cin;
using std::endl;
using std::set;

struct TNode {
    int data;
    TNode* lchild;
    TNode* rchild;
};

struct TNodeCmp {
    bool operator()(const TNode* lhs, const TNode* rhs) const {
        return lhs->data <= rhs->data;
    }
};

int CalculateWPL(TNode* root, int level) {
    if (root->lchild == nullptr && root->rchild == nullptr)
        return root->data * level;
    int left = 0;
    int right = 0;
    left = CalculateWPL(root->lchild, level + 1);
    right = CalculateWPL(root->rchild, level + 1);
    return left + right;
}

int main() {
    int n;
    cin >> n;
    set<TNode*, TNodeCmp> s;
    // 先将所有节点存储到集合中
    // 进行自动排序
    for (int i = 0; i < n; ++i) {
        TNode* node = new TNode();
        cin >> node->data;
        node->lchild = nullptr;
        node->rchild = nullptr;
        s.insert(node);
    }

    // 开始建树
    while (s.size() > 1) {
        // 先取出最小的两个节点
        set<TNode*, TNodeCmp>::iterator fst = s.begin(); // first element
        set<TNode*, TNodeCmp>::iterator sec = s.begin(); // second element
        sec++;
        TNode* newNode = new TNode();
        newNode->data = (*fst)->data + (*sec)->data;
        newNode->lchild = *fst;
        newNode->rchild =  *sec;
        // 删除两个最小值节点
        s.erase(fst);
        s.erase(sec);
        // 将新节点插入
        s.insert(newNode);
    } // 循环结束后,s中仅剩一个根节点
    // 哈夫曼树构造完成
    // 计算哈夫曼树的WPL
    int wpl = CalculateWPL(*s.begin(), 0);
    cout << wpl << endl;

    return 0;
}