def segment_by_class(data: Union[np.ndarray, list[int]], classes: np.ndarray, num_classes: int) -> list[np.ndarray]:
    """
    Returns segmented data according to specified features.

    :param data: Data to be segmented.
    :param classes: Classes used to segment data, e.g., segment according to predicted label or to `y_train` or other
                    array of one hot encodings the same length as data.
    :param num_classes: How many features.
    :return: Segmented data according to specified features.
    """
    by_class: list[list[int]] = [[] for _ in range(num_classes)]
    for indx, feature in enumerate(classes):
        if len(classes.shape) == 2 and classes.shape[1] > 1:
            assigned = int(np.argmax(feature))
        else:
            assigned = int(feature)
        by_class[assigned].append(data[indx])

    return [np.asarray(i) for i in by_class]

    Returns segmented data according to specified features.

    :param data: Data to be segmented.
    :param classes: Classes used to segment data, e.g., segment according to predicted label or to `y_train` or other
                    array of one hot encodings the same length as data.
    :param num_classes: How many features.
    :return: Segmented data according to specified features.
    
art/utils.py