diff --git a/lab2/assignment1.R b/lab2/assignment1.R
index ea564c33280ec6f09d0bc3b9bb929783ab84aef4..9b78e10a34baab9fb90e440dc8e315a66d988b8c 100644
--- a/lab2/assignment1.R
+++ b/lab2/assignment1.R
@@ -34,40 +34,37 @@ print(paste("MSE on the test data:", mse(test$Fat, test_pred)))
 
 #----3.----
 
-fit <- glmnet(as.matrix(X_train), y_train, alpha = 1)
 
-plot(fit, xvar = "lambda", label = TRUE)
-
-coef_matrix <- as.matrix(coef(fit))[-1, ] #ignore intercept
-
-lambda_values <- fit$lambda
+lambda_with_k_features = function(x_train,y_train,a,k = 3){
+  
+  fit_train <- glmnet(as.matrix(X_train), y_train, alpha = a )
+  
+  plot(fit_train, xvar = "lambda", label = TRUE)
+  
+  coef_matrix <- as.matrix(coef(fit_train))[-1, ] # Ignore intercept
+  
+  lambda_values <- fit_train$lambda
+  
+  num_non_zero <- apply(coef_matrix != 0, 2, sum) # Sum of num of cols wit non zero coeff, i.e. number of features
+  
+  lambda_values[num_non_zero == k] # lambda value wit 3 non zero coeff, i.e. number of features
 
-num_non_zero <- apply(coef_matrix != 0, 2, sum)
+}
 
-lambda_with_3_features <- lambda_values[num_non_zero == 3]
+lambda_with_3_features = lambda_with_k_features(x_train,y_train, 1,3)
 
 print(lambda_with_3_features)
 
 
 #----4.----
 
-fit <- glmnet(as.matrix(X_train), y_train, alpha = 0)
-
-plot(fit, xvar = "lambda", label = TRUE)
-
-coef_matrix <- as.matrix(coef(fit))[-1, ] #ignore intercept
-
-lambda_values <- fit$lambda
-
-num_non_zero <- apply(coef_matrix != 0, 2, sum)
-
-lambda_with_3_features <- lambda_values[num_non_zero == 3]
+lambda_with_3_features = lambda_with_k_features(x_train, y_train,0,3)
 
 print(lambda_with_3_features)
 
 #----5.----
 
-fit <- cv.glmnet(as.matrix(X_train), y_train, alpha = 1)
+fit <- cv.glmnet(as.matrix(X_train), y_train, alpha = 1) # Cross-validation
 
 plot(fit, xvar = "lambda", label = TRUE)
 
@@ -77,4 +74,4 @@ coef(fit, s="lambda.min")
 
 y_hat = predict(fit, newx=as.matrix(X_train), s="lambda.min")
 
-plot(y_train, y_hat)
+plot(y_train, y_hat, xlab = "Observed test values", ylab = "Predicted test values", main = "Original versus predicted test values ", col = "blue")