The Classification step of the CEM algorithm

def classify(self, table, weighted_choice=False, transform=None):
        """ The Classification step of the CEM algorithm """
        assert table.shape[1] == self.numgrp
        if weighted_choice:
            if transform is not None:
                probs = transform_fn(table.copy(), transform)  #
            else:
                probs = table.copy()
            cmprobs = probs.cumsum(1)
            logger.info('Probabilities\n{}'.format(probs))
            r = np.random.random(cmprobs.shape[0])
            search = np.apply_along_axis(np.searchsorted, 1, cmprobs, r) # Not very efficient
            assignment = np.diag(search)
        else:
            probs = table
            assignment = np.where(probs==probs.max(1)[:, np.newaxis])[1]
        logger.info('Assignment\n{}'.format(assignment))
        assignment = self._fill_empty_groups(probs, assignment)  # don't want empty groups
        new_partition = Partition(tuple(assignment))
        self.set_partition(new_partition)