Quick method:

    One of the biggest challenges for classification models is an imbalance of
    classes in the training data. This function vizualizes the relationship of
    the support for each class in both the training and test data by
    displaying how frequently each class occurs as a bar graph.

    The figure can be displayed in two modes:

    1. Balance mode: show the frequency of each class in the dataset.
    2. Compare mode: show the relationship of support in train and test data.

    Balance mode is the default if only y_train is specified. Compare mode
    happens when both y_train and y_test are specified.

    Parameters
    ----------
    y_train : array-like
        Array or list of shape (n,) that containes discrete data.

    y_test : array-like, optional
        Array or list of shape (m,) that contains discrete data. If
        specified, the bar chart will be drawn in compare mode.

    ax : matplotlib Axes, default: None
        The axis to plot the figure on. If None is passed in the current axes
        will be used (or generated if required).

    labels: list, optional
        A list of class names for the x-axis if the target is already encoded.
        Ensure that the labels are ordered lexicographically with respect to
        the values in the target. A common use case is to pass
        LabelEncoder.classes\_ as this parameter. If not specified, the labels
        in the data will be used.

    kwargs: dict, optional
        Keyword arguments passed to the super class. Here, used
        to colorize the bars in the histogram.

    Returns
    -------
    ax : matplotlib axes
        Returns the axes that the class balance plot was drawn on.

def class_balance(y_train, y_test=None, ax=None, labels=None, **kwargs):
    """Quick method:

    One of the biggest challenges for classification models is an imbalance of
    classes in the training data. This function vizualizes the relationship of
    the support for each class in both the training and test data by
    displaying how frequently each class occurs as a bar graph.

    The figure can be displayed in two modes:

    1. Balance mode: show the frequency of each class in the dataset.
    2. Compare mode: show the relationship of support in train and test data.

    Balance mode is the default if only y_train is specified. Compare mode
    happens when both y_train and y_test are specified.

    Parameters
    ----------
    y_train : array-like
        Array or list of shape (n,) that containes discrete data.

    y_test : array-like, optional
        Array or list of shape (m,) that contains discrete data. If
        specified, the bar chart will be drawn in compare mode.

    ax : matplotlib Axes, default: None
        The axis to plot the figure on. If None is passed in the current axes
        will be used (or generated if required).

    labels: list, optional
        A list of class names for the x-axis if the target is already encoded.
        Ensure that the labels are ordered lexicographically with respect to
        the values in the target. A common use case is to pass
        LabelEncoder.classes\_ as this parameter. If not specified, the labels
        in the data will be used.

    kwargs: dict, optional
        Keyword arguments passed to the super class. Here, used
        to colorize the bars in the histogram.

    Returns
    -------
    ax : matplotlib axes
        Returns the axes that the class balance plot was drawn on.
    """
    # Instantiate the visualizer
    visualizer = ClassBalance(ax=ax, labels=labels, **kwargs)

    # Fit and transform the visualizer (calls draw)
    visualizer.fit(y_train, y_test)
    visualizer.finalize()

    # Return the axes object on the visualizer
    return visualizer.ax