updated the code 

updated the code for the new grid 
used junit 4.5