1 files changed, 44 insertions, 1 deletions
diff --git a/rst_parser/rst.cc b/rst_parser/rst.cc
index f6b295b3..c4ce898e 100644
--- a/rst_parser/rst.cc
+++ b/rst_parser/rst.cc
@@ -2,6 +2,49 @@
 
 using namespace std;
 
-StochasticForest::StochasticForest(const ArcFactoredForest& af) {
+// David B. Wilson. Generating Random Spanning Trees More Quickly than the Cover Time.
+
+TreeSampler::TreeSampler(const ArcFactoredForest& af) : forest(af), usucc(af.size() + 1) {
+  // edges are directed from modifiers to heads, to the root
+  for (int m = 1; m <= forest.size(); ++m) {
+    SampleSet<double>& ss = usucc[m];
+    for (int h = 0; h <= forest.size(); ++h)
+      ss.add(forest(h-1,m-1).edge_prob.as_float());
+  }
 }
 
+void TreeSampler::SampleRandomSpanningTree(EdgeSubset* tree) {
+  MT19937 rng;
+  const int r = 0;
+  bool success = false;
+  while (!success) {
+    int roots = 0;
+    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) {
+      int u = i;
+      if (in_tree[u]) continue;
+      while(!in_tree[u]) {
+        next[u] = rng.SelectSample(usucc[u]);
+        u = next[u];
+      }
+      u = i;
+      cerr << (u-1);
+      while(!in_tree[u]) {
+        in_tree[u] = true;
+        u = next[u];
+        cerr << " > " << (u-1);
+        if (u == r) { ++roots; }
+      }
+      cerr << endl;
+    }
+    assert(roots > 0);
+    if (roots > 1) {
+      cerr << "FAILURE\n";
+    } else {
+      success = true;
+    }
+  }
+};
+