input: Write a python function to find the maximum element in an array.
  assert find_max([1, 2, 3, 4, 5]) == 5
  assert find_max([-10, -5, 0, 5, 10]) == 10
  assert find_max([100, 50, 200, 150]) == 200
output: def find_max(arr):
      if not arr:
          return None
      max_element = arr[0]
      for num in arr:
          if num > max_element:
              max_element = num
      return max_element