题解 | 搬水果

#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) {
    // 哈夫曼树所有节点只有 度为0 和度为 2两种可能
    if (root->lchild == nullptr && root->rchild == nullptr)
        return root->data * level;
    int left = CalculateWPL(root->lchild, level + 1);
    int right = CalculateWPL(root->rchild, level + 1);
    return left + right;
}

int main() {
    int n;
    cin >> n;
    set<TNode*, TNodeCmp> s;
    // 先将所有输入的数据变为节点
    // 然后放进s中,进行自动排序
    for (int i = 0; i < n; ++i) {
        TNode* node = new TNode();
        cin >> node->data;
        node->lchild = nullptr;
        node->rchild = nullptr;
        s.insert(node);
    }

    // 根据s中的数据开始构建哈夫曼树
    while (s.size() > 1) {
        set<TNode*, TNodeCmp>::iterator fst = s.begin();
        set<TNode*, TNodeCmp>::iterator sec = s.begin();
        sec++;  // 让sec指向s中的第二小的元素
        TNode* newNode = new TNode();
        newNode->data = (*fst)->data + (*sec)->data;
        newNode->lchild = *fst;
        newNode->rchild = *sec;
        // 删除s中的两个最小元素
        // 将合并后的子树的根节点放入s
        s.erase(fst);
        s.erase(sec);
        s.insert(newNode);
    }// 循环结束之后,s中只剩下哈夫曼树的根节点

    // 根据哈夫曼树求wpl
    int wpl = CalculateWPL(*s.begin(), 0);
    cout << wpl << endl;
    return 0;
}