This could be done with two less line with `label_binarize `.