updated code from competition ;)