diff --git a/05. DiagonalDifference/solution.py b/05. DiagonalDifference/solution.py
index ef676bd..a88abba 100644
--- a/05. DiagonalDifference/solution.py	
+++ b/05. DiagonalDifference/solution.py	
@@ -6,14 +6,6 @@
     Url : https://www.hackerrank.com/challenges/diagonal-difference/problem
 '''
 #!/bin/python3
-
-import math
-import os
-import random
-import re
-import sys
-
-#
 # Complete the 'diagonalDifference' function below.
 #
 # The function is expected to return an INTEGER.
@@ -22,22 +14,12 @@
 
 def diagonalDifference(arr):
     # Write your code here
+    n = len(arr)
     d1 = sum(arr[i][i] for i in range(n))
     d2 = sum(arr[i][n-i-1] for i in range(n))
     return abs(d1 - d2)
-    
-if __name__ == '__main__':
-    fptr = open(os.environ['OUTPUT_PATH'], 'w')
-
-    n = int(input().strip())
-
-    arr = []
-
-    for _ in range(n):
-        arr.append(list(map(int, input().rstrip().split())))
-
-    result = diagonalDifference(arr)
 
-    fptr.write(str(result) + '\n')
 
-    fptr.close()
+assert diagonalDifference([[11,2,4], [4,5,6], [10,8,-12]]) == 15
+assert diagonalDifference([[1,2,3], [4,5,6], [9,8,9]]) == 2
+assert diagonalDifference([[1,1,1,1], [1,1,1,1], [1,1,1,1], [1,1,1,1]]) == 0
diff --git a/06. PlusMinus/solution.py b/06. PlusMinus/solution.py
index 7979b16..2c331e4 100644
--- a/06. PlusMinus/solution.py	
+++ b/06. PlusMinus/solution.py	
@@ -9,29 +9,19 @@
 '''
 #!/bin/python3
 
-import math
-import os
-import random
-import re
-import sys
-
 # Complete the plusMinus function below.
 def plusMinus(arr):
     countPositive = 0
     countNegative =0
+    n = len(arr)
     for i in range(n):
         if arr[i] > 0:
             countPositive += 1 
         elif arr[i] < 0:
             countNegative += 1
     countZero = n - countPositive-countNegative
-    print(countPositive/n)
-    print(countNegative/n)
-    print(countZero/n)
+    print("Count Positive:", countPositive/n)
+    print("Count Negative:", countNegative/n)
+    print("Count Zero:", countZero/n)
     
-if __name__ == '__main__':
-    n = int(input())
-
-    arr = list(map(int, input().rstrip().split()))
-
-    plusMinus(arr)
+plusMinus([-4, 3, -9, 0, 4, 1])