Swapped data structure and cleaned up a bit