Update prime_number.py

Author: saudzahirr

prime_number.py

@@ -4,14 +4,23 @@
 # Q 17. Prime Number.
 
 def prime(N):
-    flag = 0
-    for i in range(2, N):
-        if N%i == 0:
-            flag = 1
-    if flag == 1:
-        print("N is not a Prime Number.")
+    if N == 1:
+        print("One is neither Prime nor Composite.")
+    
+    elif N == 0:
+        print("Zero is neither Prime nor Composite.")
+    
     else:
-        print("N is a Prime Number.")
+        flag = 0
+        for i in range(2, N):
+            if N%i == 0:
+                flag = 1
+        
+        if flag == 1:
+            print("N = ", N, "is not a Prime Number.")
+        
+        else:
+            print("N = ", N, "is a Prime Number.")
 
 
 print("Enter the number...")