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
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112

#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;
  vector<vector<bool>> 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 < 8; ++i) v.push_back(1 & (n->c >> i));
    }
    else {
      v.push_back(0);
      if (n->l) {
        v.push_back(1);
        encode_tree(n->l, v);
      }
      else v.push_back(0);
      if (n->r) {
        v.push_back(1);
        encode_tree(n->r, v);
      }
      else v.push_back(0);
    }
  }

  int idx = 0;
  void decode_tree(node * n, vector<bool> & v) {
    if (v[idx++] == 1) {
      for (int i = 0; i < 8; ++i) n->c |= (1 << v[idx++]);
    }
    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 = 0; c < 128; ++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 (n->c) {
      s += n->c;
      //cout << (int)n->c << '\n';
      if (idx > v.size()) return;
      solve(root, v, s);
    }
    else {
      if (v[idx++] == 0) solve(n->l, v, s);
      else 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;
  }
}