Updating Write_a_function.py

[arsh-purewal]

Your code was not giving me correct output for the test cases on hackerrank, for accurate results you can replace your code with mine as it is less lengthy and giving me accurate results.