Fix random heuristic

1 files changed, 20 insertions, 6 deletions

diff --git a/src/random_heuristic.c b/src/random_heuristic.c
index a2d0747..0076728 100644
--- a/src/random_heuristic.c
+++ b/src/random_heuristic.c
@@ -100,13 +100,13 @@ struct online_return random_heuristic_online(struct data * data, struct front_it
   }
   long int tree_size = 0;
   for(struct tree_item * tree_node = tree; tree_node != NULL; tree_node = tree_node->next, tree_size++);
+  
+  /* Choose tree node randomly */
+  long int position = rand() % tree_size;
+  struct tree_item * node = tree;
+  for(int i = 0; i < position; i++,node=node->next);
 
   do {
-    /* Choose tree node randomly */
-    long int position = rand() % tree_size;
-    struct tree_item * node = tree;
-    for(int i = 0; i < position; i++,node=node->next);
-    
     /* Check if node is leaf */
     if(node->depth == data->N) {
       front_last->next = new_front_item(node->values.p, node->values.w, front_last, NULL);
@@ -127,9 +127,16 @@ struct online_return random_heuristic_online(struct data * data, struct front_it
       /* Reduce tree_size */
       tree_size--;
 
+      if(tree != NULL) {
+        /* Select new random node */
+        position = rand() % tree_size;
+        node = tree;
+        for(int i = 0; i < position; i++,node=node->next);
+      }
+
       continue;
     }
-    
+
     /* Make left node (select item) */
     struct tree_item * left = new_tree_item(
       node->depth + 1,
@@ -165,6 +172,13 @@ struct online_return random_heuristic_online(struct data * data, struct front_it
     }
     free(old);
 
+    /* Choose left or right randomly for node */
+    if(rand() % 2 == 0) {
+      node = left;
+    } else {
+      node = right;
+    }
+
     /* Increment tree_size */
     tree_size++;
   } while(tree != NULL);