Add edge cases for equilibrium index

Author: NeuroByte79

data_structures/arrays/equilibrium_index_in_array.py

@@ -1,52 +1,87 @@
+
 """
 Find the Equilibrium Index of an Array.
-Reference: https://www.geeksforgeeks.org/equilibrium-index-of-an-array/
 
-Python doctest can be run with the following command:
+Reference:
+https://www.geeksforgeeks.org/equilibrium-index-of-an-array/
+
+Python doctest can be run with:
+
 python -m doctest -v equilibrium_index_in_array.py
 
-Given a sequence arr[] of size n, this function returns
-an equilibrium index (if any) or -1 if no equilibrium index exists.
+Given an array arr of size n, return an equilibrium index
+if one exists, otherwise return -1.
 
-The equilibrium index of an array is an index such that the sum of
-elements at lower indexes is equal to the sum of elements at higher indexes.
+An equilibrium index is an index where the sum of all
+elements to the left equals the sum of all elements
+to the right.
 
+Examples:
 
+    >>> equilibrium_index([-7, 1, 5, 2, -4, 3, 0])
+    3
 
-Example Input:
-arr = [-7, 1, 5, 2, -4, 3, 0]
-Output: 3
+    >>> equilibrium_index([1, 2, 3, 4, 5])
+    -1
 
+    >>> equilibrium_index([1, 1, 1, 1, 1])
+    2
+
+    >>> equilibrium_index([2, 4, 6, 8, 10, 3])
+    -1
 """
 
 
 def equilibrium_index(arr: list[int]) -> int:
     """
-    Find the equilibrium index of an array.
+    Find the first equilibrium index of an array.
 
     Args:
-        arr (list[int]): The input array of integers.
+        arr: The input array of integers.
 
     Returns:
-        int: The equilibrium index or -1 if no equilibrium index exists.
+        The first equilibrium index, or -1 if none exists.
 
     Examples:
+        >>> equilibrium_index([])
+        -1
+
+        >>> equilibrium_index([5])
+        0
+
         >>> equilibrium_index([-7, 1, 5, 2, -4, 3, 0])
         3
+
         >>> equilibrium_index([1, 2, 3, 4, 5])
         -1
+
         >>> equilibrium_index([1, 1, 1, 1, 1])
         2
-        >>> equilibrium_index([2, 4, 6, 8, 10, 3])
-        -1
+
+        >>> equilibrium_index([0, 0, 0])
+        0
+
+        >>> equilibrium_index([-1, -1, -1])
+        1
+
+        >>> equilibrium_index([1, -1, 0])
+        2
+
+    Time Complexity:
+        O(n), where n is the length of the array.
+
+    Space Complexity:
+        O(1), using only constant extra space.
     """
     total_sum = sum(arr)
     left_sum = 0
 
     for i, value in enumerate(arr):
         total_sum -= value
+
         if left_sum == total_sum:
             return i
+
         left_sum += value
 
     return -1
@@ -55,4 +90,4 @@ def equilibrium_index(arr: list[int]) -> int:
 if __name__ == "__main__":
     import doctest
 
-    doctest.testmod()
+    doctest.testmod()