Improved task 1 solution

Signed-off-by: Nikola Iliev <nikola.iliev93@gmail.com>