input: Write a function to find the maximum element in the given array. If the array is empty, return None. The function should be called `find_max`.
  
  assert find_max([1, 10, 50, 3, 8])==50
  assert find_max([5, 7, 9, 3, 2])==9
  assert find_max([-1, -5, -3, -8])==-1
  assert find_max([])==None
output: def find_max(arr):
      if not arr:
          return None
      return max(arr)