diff --git a/assignment-2.py b/assignment-2.py
index 956b815..c8187f5 100644
--- a/assignment-2.py
+++ b/assignment-2.py
@@ -1,92 +1,85 @@
-# DSL - Assignment 2
+# List for storing marks
+marks = []
+total = int(input("Total number of students are:\t"))
 
-students=[]
-marks=[]
-total_students_in=int(input("Total number of students are: \t"))
+# Input marks
+def marksInput():
+    print("\n----------\nNOTE: PLEASE ENTER MARKS OUT OF 50. ENTER '-1' FOR ABSENT STUDENTS.\n----------\n")
+    for i in range(total):
+        enterMarks = int(input(f"Enter marks for student {i+1}:\t"))
+        marks.append(enterMarks)
+    print(f"\n-----\nMarks of {total} students are:\t{marks}\n-----")
 
-# Entering marks
-def total_students_fn():
-	for i in range(total_students_in):
-		marks_in=int(input("Enter marks for DSA subject for each student (out of 50) (enter -1 for absent students): \t"))
-		marks.append(marks_in)
-	print("Marks of", total_students_in, "students in DSA subject are:", marks)
-
-total_students_fn()
-
-# Option 1 = Average
+# Option 1 = Average score
 def average_marks():
-	for i in marks:
-		if i<0:
-			break
-		else:
-			average_calc=sum(marks)/total_students_in
-	print("Average score of", total_students_in, "is: \t", average_calc)
+    marks_withoutAbsent = []
+    for i in marks:
+        if (i >= 0):
+            marks_withoutAbsent.append(i)
+        else:
+            break
+    average_calc = sum(marks_withoutAbsent) / total
+    print(f"\n------\nAverage score of {total} students is:\t{average_calc}\n-----")
 
-# Option 2 = Highest and lowest
+# Option 2 = High and low marks
 def high_low():
-   mini=marks[0]
-   maxi=marks[0]
-   for i in range (len(marks)):
-      if (maxi<marks[i] and marks[i]>-1):
-         maxi=marks[i]
-   for j in range (len(marks)):
-      if (mini>marks[j] and marks[j]>-1):
-         mini=marks[j]
-   print("Highest marks are: \t", maxi)
-   print("Lowest marks are: \t", mini)
-   
-# Option 3 = Absent
+    maxi = marks[0] # Initialize maxi with the first element of the marks list
+    mini = marks[0] # Initialize mini with the first element of the marks list
+    for i in range(len(marks)):
+        if (maxi < marks[i] and marks[i] > -1):
+            maxi = marks[i]
+    for j in range(len(marks)):
+        if (mini > marks[j] and marks[j] > -1):
+            mini = marks[j]
+    print(f"\n-----\nHighest score is:\t{maxi}\nLowest score is:\t{mini}\n-----")
+
+# Option 3 = Absent count
 def absent():
-	absent_count=0
-	for i in marks:
-		if i==-1:
-			absent_count+=1	
-	print("Total number of absent students out of", total_students_in, "are: \t", absent_count)
+    absent_count = 0;
+    for i in marks:
+        if (i < 0):
+            absent_count+=1
+        else:
+            continue
+    print(f"\n-----\nTotal absent students are:\t{absent_count}\n-----")
 
-# Option 4 = Highest freq
+# Option 4 = Highest frequency
 def high_freq():
-   freq_count=0
-   for i in range(len(marks)):
-      if (i>=0):
-         count=0
-         for j in range(len(marks)):
-            if (marks[i]==marks[j]):
-               count+=1
-               if (freq_count<count):
-                  freq_count=count
-      else:
-         break
-   print("Highest frequency is: \t", freq_count)
+    freq_count = 0 # Initialise frequency counter
+    for i in range(len(marks)): # Iterate through marks list using i
+        if (marks[i] >= 0): # Only consider non-negative marks (since -1 is absent)
+            temp_count = 0 # Initialise temporary counter
+            for j in range(len(marks)): # Iterate through marks list using j again
+                if (marks[i] == marks[j]):
+                    temp_count+=1
+                    if (freq_count < temp_count): # If temp_count is greater than freq_count value, then
+                        freq_count = temp_count # Make freq_count equal to temp_count
+    print(f"\n-----\nHighest frequency is:\t{freq_count}\n-----")
 
-# Choosing an option
-def choose_optn():
-   while True:
-      print("Choose an operation to perform:")
-      print("1. Average score of all students")
-      print("2. Highest and lowest score out of all students")
-      print("3. Count of students absent for test")
-      print("4. Display marks with highest frequency")
-      print("5. Exit")
-      optn=int(input("Enter option number (1-5): \t"))
-      if optn==1:
-         average_marks()
-         print("Designed and Engineered by Kshitij")
-      elif optn==2:
-         high_low()
-         print("Designed and Engineered by Kshitij")
-      elif optn==3:
-         absent()
-         print("Designed and Engineered by Kshitij")
-      elif optn==4:
-         high_freq()
-         print("Designed and Engineered by Kshitij")
-      elif optn==5:
-         print("## Designed and Engineered by Kshitij\n## END OF CODE.")
-         quit()
-      else:
-         print("\n Invalid option selected. Please choose from range 1 to 5. \n")
-         choose_optn()
-
-choose_optn()
-
-## END OF CODE
+def main():
+    marksInput()
+    
+    while (1):
+        print("\n----- MAIN MENU -----")
+        print("1 -> Average score of class")
+        print("2 -> Highest and lowest score of class")
+        print("3 -> Total absent students")
+        print("4 -> Marks with highest frequency")
+        print("5 -> Exit")
+        
+        optn = int(input("Choose an option (1-5):\t"))
+        
+        if (optn == 1):
+            average_marks()
+        elif (optn == 2):
+            high_low()
+        elif (optn == 3):
+            absent()
+        elif (optn == 4):
+            high_freq()
+        elif (optn == 5):
+            print("\n\n## END OF CODE\n\n")
+            quit()
+        else:
+            print("\n\nPlease choose a valid option (1-5)\n\n")
+main()