diff --git a/lab2/assignment1.R b/lab2/assignment1.R
index b706dd3ddd11555b28535a73fd93e24c35d868ed..2df7f33bd40b32c4aabc28ec884e98905b1aec29 100644
--- a/lab2/assignment1.R
+++ b/lab2/assignment1.R
@@ -34,8 +34,15 @@ print(paste("MSE on the test data:", mse(test$Fat, test_pred)))
 #----3.----
 
 
-lambda_with_k_features = function(x_train, y_train, alpha, k = 3) {
-  fit_train <- glmnet(as.matrix(X_train), y_train, alpha = alpha)
+lambda_with_k_features = function(x_train,
+                                  y_train,
+                                  alpha,
+                                  k = 3,
+                                  nlambda = 1000) {
+  fit_train <- glmnet(as.matrix(X_train),
+                      y_train,
+                      alpha = alpha,
+                      nlambda = nlambda)
   
   plot(fit_train, xvar = "lambda", label = TRUE)
   grid()
@@ -48,7 +55,16 @@ lambda_with_k_features = function(x_train, y_train, alpha, k = 3) {
   num_non_zero <- apply(coef_matrix != 0, 2, sum)
   
   # lambda value with k non zero coeff, i.e. number of features
-  lambda_values[num_non_zero == k]
+  lambda_values <- lambda_values[num_non_zero == k]
+  
+  if (length(lambda_values) == 0) {
+    return()
+  }
+  
+  return(list(
+    min = min(lambda_values),
+    max = max(lambda_values)
+  ))
   
 }