1 files changed, 44 insertions, 12 deletions

diff --git a/rst_parser/rst.cc b/rst_parser/rst.cc
index c4ce898e..bc91330b 100644
--- a/rst_parser/rst.cc
+++ b/rst_parser/rst.cc
@@ -3,45 +3,77 @@
 using namespace std;
 
 // David B. Wilson. Generating Random Spanning Trees More Quickly than the Cover Time.
-
+// this is an awesome algorithm
 TreeSampler::TreeSampler(const ArcFactoredForest& af) : forest(af), usucc(af.size() + 1) {
-  // edges are directed from modifiers to heads, to the root
+  // edges are directed from modifiers to heads, and finally to the root
+  vector<double> p;
   for (int m = 1; m <= forest.size(); ++m) {
+#if USE_ALIAS_SAMPLER
+    p.clear();
+#else
     SampleSet<double>& ss = usucc[m];
-    for (int h = 0; h <= forest.size(); ++h)
-      ss.add(forest(h-1,m-1).edge_prob.as_float());
+#endif
+    double z = 0;
+    for (int h = 0; h <= forest.size(); ++h) {
+      double u = forest(h-1,m-1).edge_prob.as_float();
+      z += u;
+#if USE_ALIAS_SAMPLER
+      p.push_back(u);
+#else
+      ss.add(u);
+#endif
+    }
+#if USE_ALIAS_SAMPLER
+    for (int i = 0; i < p.size(); ++i) { p[i] /= z; }
+    usucc[m].Init(p);
+#endif
   }
 }
 
-void TreeSampler::SampleRandomSpanningTree(EdgeSubset* tree) {
-  MT19937 rng;
+void TreeSampler::SampleRandomSpanningTree(EdgeSubset* tree, MT19937* prng) {
+  MT19937& rng = *prng;
   const int r = 0;
   bool success = false;
   while (!success) {
     int roots = 0;
+    tree->h_m_pairs.clear();
+    tree->roots.clear();
     vector<int> next(forest.size() + 1, -1);
     vector<char> in_tree(forest.size() + 1, 0);
     in_tree[r] = 1;
-    for (int i = 0; i < forest.size(); ++i) {
+    //cerr << "Forest size: " << forest.size() << endl;
+    for (int i = 0; i <= forest.size(); ++i) {
+      //cerr << "Sampling starting at u=" << i << endl;
       int u = i;
       if (in_tree[u]) continue;
       while(!in_tree[u]) {
+#if USE_ALIAS_SAMPLER
+        next[u] = usucc[u].Draw(rng);
+#else
         next[u] = rng.SelectSample(usucc[u]);
+#endif
         u = next[u];
       }
       u = i;
-      cerr << (u-1);
+      //cerr << (u-1);
+      int prev = u-1;
       while(!in_tree[u]) {
         in_tree[u] = true;
         u = next[u];
-        cerr << " > " << (u-1);
-        if (u == r) { ++roots; }
+        //cerr << " > " << (u-1);
+        if (u == r) {
+          ++roots;
+          tree->roots.push_back(prev);
+        } else {
+          tree->h_m_pairs.push_back(make_pair<short,short>(u-1,prev));
+        }
+        prev = u-1;
       }
-      cerr << endl;
+      //cerr << endl;
     }
     assert(roots > 0);
     if (roots > 1) {
-      cerr << "FAILURE\n";
+      //cerr << "FAILURE\n";
     } else {
       success = true;
     }