diff --git a/R/NetworkCausalTree.R b/R/NetworkCausalTree.R
index fa04db6..08a91fa 100644
--- a/R/NetworkCausalTree.R
+++ b/R/NetworkCausalTree.R
@@ -157,17 +157,29 @@ NetworkCausalTree <- function(X, Y, W,
   tree$FILTER <- as.character(tree$FILTER)
   tree$FILTER[is.na(tree$FILTER)] <- "NA"
   
+  clusters_est <- setdiff(1:k, clusters_disc)
+  est_units <- which(K %in% clusters_est)
+  
+  A_est <- A[est_units, est_units]
+  Ne_est <- rowSums(A_est)
+  Ne_treated_est <- as.vector(A_est %*% W[est_units])
+  G_est <- rep(1, length(est_units))
+  G_est[Ne_treated_est == 0] <- 0
+  Ne_list_est <- lapply(seq_along(est_units), function(i) {
+    which(A_est[i, ] > 0)
+  })
+  
   # Compute estimates on the full tree
   tree_with_estimates <- compute_effects_nct(output = output,
                                              nct_partition = tree,
-                                             N = N,
-                                             W = W,
-                                             G = G,
-                                             Y = Y,
-                                             X = X,
-                                             p = p,
-                                             Ne = Ne,
-                                             Ne_list = Ne_list,
+                                             N = length(est_units),
+                                             W = W[est_units],
+                                             G = G_est,
+                                             Y = Y[est_units],
+                                             X = X[est_units, , drop = FALSE],
+                                             p = p[est_units],
+                                             Ne = Ne_est,
+                                             Ne_list = Ne_list_est,
                                              minsize = minsize)
   
   return(tree_with_estimates)
diff --git a/paper/images/NCT_algorithm.png b/paper/images/NCT_algorithm.png
deleted file mode 100644
index 0d2d449..0000000
Binary files a/paper/images/NCT_algorithm.png and /dev/null differ
diff --git a/tests/testthat/test-plot_NCT.R b/tests/testthat/test-plot_NCT.R
index 282e3c4..48e2dcc 100644
--- a/tests/testthat/test-plot_NCT.R
+++ b/tests/testthat/test-plot_NCT.R
@@ -2,11 +2,13 @@ test_that("plot_NCT runs on a real tiny NCT object", {
   skip_on_cran()
   
   set.seed(100)
-  d <- data_generator_direct_indirect(N = 100, k = 5, M = 2, p = rep(0.3,100), remove_isolates = TRUE)
+  d <- data_generator_direct_indirect(N = 500, k = 10, M = 2, p = rep(0.3, 500), remove_isolates = TRUE)
   
   fit <- NetworkCausalTree(
     Y = d$Y, W = d$W, X = d$X, A = d$A, p = d$p,
-    method = "OF", depth = 1, minsize = 20,
+    K = d$K,
+    effect_weights = c(1,0,0,0),
+    method = "singular", depth = 1, minsize = 20,
     ratio_disc = 0.5
   )
   
@@ -22,11 +24,13 @@ test_that("plot_NCT works in detection mode", {
   skip_on_cran()
   
   set.seed(100)
-  d <- data_generator_direct_indirect(N = 100, k = 5, M = 2, p = rep(0.3,100), remove_isolates = TRUE)
+  d <- data_generator_direct_indirect(N = 500, k = 10, M = 2, p = rep(0.3, 500), remove_isolates = TRUE)
   
   fit <- NetworkCausalTree(
     Y = d$Y, W = d$W, X = d$X, A = d$A, p = d$p,
-    method = "OF", depth = 1, minsize = 20,
+    K = d$K,
+    effect_weights = c(1,0,0,0),
+    method = "singular", depth = 1, minsize = 20,
     ratio_disc = 0.5,
     output = "detection"
   )