Update main.go

Shorten code and add header.