【Openjudge】ハフマン符号化ツリー(STL)

今回はSTLを使いました

#include
#include

using namespace std;

struct _node{
	int weight;
	_node* left;
	_node* right;
	_node():weight(0), left(NULL), right(NULL){};
	friend bool operator>(_node a, _node b){
		if (a.weight > b.weight)
			return true;
		return false;
	}
	void merge(_node &lc, _node &rc){
		weight = lc.weight + rc.weight;
		left = &lc;
		right = &rc;
	}
};

void travel_and_sum(_node *root, int &sum, int &depth){
	if (root->left == NULL && root->right == NULL){
		sum += root->weight * depth;
		depth --;
		return;
	}
	if (root->left){
		depth ++;
		travel_and_sum(root->left, sum, depth);
		depth ++;
		travel_and_sum(root->right, sum, depth);
	}
	depth --;
	return;
}

void swap(_node &a, _node &b){
	_node t = a;
	a = b;
	b = t;
}

int main(){
	int amount, current_size;
	cin >> amount;
	current_size = amount;
	_node *heap_array = new _node[amount];
	for (int i = 0; i < amount; i ++){
		int t;
		cin >> t;
		heap_array[i].weight = t;
	}
	make_heap(heap_array,heap_array + amount, greater<_node>());

	_node *parent, *leftchild, *rightchild;
	for (int i = 0;i < amount - 1; i ++){
		leftchild = new _node(heap_array[0]);
		pop_heap(heap_array, heap_array + current_size --, greater<_node>());
		rightchild = new _node(heap_array[0]);
		pop_heap(heap_array, heap_array + current_size --, greater<_node>());
		parent = new _node;
		parent->merge(*leftchild, *rightchild);
		heap_array[current_size] = *parent;
		push_heap(heap_array, heap_array + current_size ++, greater<_node>());
	}
	int sum = 0, depth = 0;
	travel_and_sum(parent, sum, depth);
	cout << sum << endl;
}