modified main.cpp for better output