1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79

#include "huffman.h"
#include "common.h"

namespace huffman {
  struct node {
    char c; int f; node * l, * r;
    node(char _c, int _f) { c = _c, f = _f, l = r = NULL; }
  };
  struct comp { bool operator()(node * a, node * b) { return a->f > b->f; } };
  node * root;
  vector<bool> path, code[128];
  void traverse(node * n) {
    if (n->l) path.push_back(0), traverse(n->l), path.pop_back();
    if (n->r) path.push_back(1), traverse(n->r), path.pop_back();
    if (n->c) code[n->c] = path;
  }

  void encode_tree(node * n, vector<bool> & v) {
    if (n->c) {
      v.push_back(1);
      for (int i = 0; i < 7; ++i) v.push_back(1 & (n->c >> i));
    }
    else {
      v.push_back(0);
      n->l ? v.push_back(1), encode_tree(n->l, v) : v.push_back(0);
      n->r ? v.push_back(1), encode_tree(n->r, v) : v.push_back(0);
    }
  }

  int idx = 0;
  void decode_tree(node * n, vector<bool> & v) {
    if (idx >= v.size()) return;
    if (v[idx++] == 1) {
      for (int i = 0; i < 7; ++i) if (v[idx++]) n->c |= (1 << i);
    }
    else {
      if (v[idx++] == 1) n->l = new node(0, 0), decode_tree(n->l, v);
      if (v[idx++] == 1) n->r = new node(0, 0), decode_tree(n->r, v);
    }    
  }

  void generate(vector<int> f) {
    priority_queue<node *, vector<node *>, comp> pq;
    for (int c = 1; c < 128; ++c) if (f[c]) {
      node * n = new node(c, f[c]); pq.push(n);
    }
    while (pq.size() > 1) {
      node * l = pq.top(); pq.pop(); node * r = pq.top(); pq.pop();
      node * n = new node(0, l->f + r->f); n->l = l, n->r = r;
      pq.push(n);
    }
    root = pq.top();
  }

  void solve(node * n, vector<bool> & v, string & s) {
	if (idx >= v.size()) return;
    if (n->c) s += n->c, solve(root, v, s);
    else v[idx++] == 0 ? solve(n->l, v, s) : solve(n->r, v, s);
  }

  vector<bool> encode(string s) {
    vector<int> f(128, 0);
    for (auto& c : s) ++f[c];
    generate(f);
    traverse(root);
    vector<bool> ret;
    encode_tree(root, ret);
    for (auto& c : s) for (auto b : code[c]) ret.push_back(b);
    return ret;
  }

  string decode(vector<bool> v) {
    root = new node(0, 0);
    decode_tree(root, v);
    string ret;
    solve(root, v, ret);
    return ret;
  }
}