fix bug in solution for #398