Update main.c

Optimisation of the display function