diff --git a/lab2/assignment2.R b/lab2/assignment2.R
index 0cd21f782752fdc403ea6cdcedfe14d36b612b27..6e4780aadc1ca5f07b7a266b6bca95ea9013cc20 100644
--- a/lab2/assignment2.R
+++ b/lab2/assignment2.R
@@ -165,15 +165,11 @@ F1(test$y, y_hat4)
 
 # ----5.----
 
-loss_matrix = matrix(c(0, 1, 5, 0), ncol = 2)
-
 fit5 <- tree(y ~ ., data = train)
 
-pruned_tree5 = prune.tree(fit5,
-                          loss = loss_matrix,
-                          best = sum(fit5$frame$var == "<leaf>"))
+probs5 <- predict(fit5, newdata = test, type = "vector")
 
-y_hat5 = predict(pruned_tree5, newdata = test, type = "class")
+y_hat5 <- ifelse(probs5[, "no"] / probs5[, "yes"] > 5, "no", "yes")
 
 table(test$y, y_hat5)