From 225f0647bda3401c9270f6056f0384ee2be4ff22 Mon Sep 17 00:00:00 2001
From: Alexander Kowarik <alex@kowarik.net>
Date: Thu, 9 Nov 2023 08:24:55 +0100
Subject: [PATCH] xgboostImpute fix for binary variables

---
 R/xgboostImpute.R | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/R/xgboostImpute.R b/R/xgboostImpute.R
index f77a070..ff0ed04 100644
--- a/R/xgboostImpute.R
+++ b/R/xgboostImpute.R
@@ -96,7 +96,11 @@ xgboostImpute <- function(formula, data, imp_var = TRUE,
       predictions <- 
         predict(mod, model.matrix(formPred,subset(data, !rhs_na & lhs_na)))
       if(currentClass=="factor"){
-        data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[predictions+1]
+        if(is.null(num_class)){
+          data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[as.numeric(predictions>.5)+1]  
+        }else{
+          data[!rhs_na & lhs_na, lhsV] <- levels(dattmp[,lhsV])[predictions+1]  
+        }
       }else{
         data[!rhs_na & lhs_na, lhsV] <- predictions  
       }