Added stuff for problem 4

Not working correctly..