Optimize remove method
Slight test modification