fixed remove method to return the correct element always