Getting the decision boundary for KNN classifier using R

Here is a tweaked approach that draws the decision boundaries as lines. I thought this would require the predicted probability for each class but after reading this answer it turns out you can just mark the predicted probability for each class as 1 wherever that class is predicted, and zero otherwise.

# Create matrices for each class where p = 1 for any point
#   where that class was predicted, 0 otherwise
n_classes = 3
class_regions = lapply(1:n_classes, function(class_num) {
    indicator = ifelse(mod.opt == class_num, 1, 0)
    mat = matrix(indicator, n.grid1, n.grid1)
})

# Set up colours
class_colors = c("#4E79A7", "#F28E2B", "#E15759")
# Add some transparency to make the fill colours less bright
fill_colors = paste0(class_colors, "60")

# Use image to plot the predicted class at each point
classes = matrix(as.numeric(mod.opt), n.grid1, n.grid1)
image(x1.grid1, x2.grid1, classes, col = fill_colors, 
      main = "plot of training data with decision boundary",
      xlab = colnames(trainxx)[1], ylab = colnames(trainxx)[2])
# Draw contours separately for each class
lapply(1:n_classes, function(class_num) {
    contour(x1.grid1, x2.grid1, class_regions[[class_num]], 
            col = class_colors[class_num],
            nlevels = TRUE, add = TRUE, lwd = 2, drawlabels = FALSE)
})
# Using pch = 21 for bordered points that stand out a bit better
points(trainxx, bg = class_colors[trainyy], 
       col = "black",
       pch = 21)

The resulting plot: