<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def cleanup_dataset(dataset, data_home=None, ext=".zip"):
 """ Removes the dataset directory and archive file from the data home directory. Parameters dataset : str The name of the dataset; should either be a folder in data home or specified in the yellowbrick.datasets.DATASETS variable. data_home : str, optional The path on disk where data is stored. If not passed in, it is looked up from YELLOWBRICK_DATA or the default returned by ``get_data_home``. ext : str, default: ".zip" The extension of the archive file. Returns ------- removed : int The number of objects removed from data_home. """

    removed = 0
    data_home = get_data_home(data_home)

    # Paths to remove
    datadir = os.path.join(data_home, dataset)
    archive = os.path.join(data_home, dataset+ext)

    # Remove directory and contents
    if os.path.exists(datadir):
        shutil.rmtree(datadir)
        removed += 1

    # Remove the archive file
    if os.path.exists(archive):
        os.remove(archive)
        removed += 1

    return removed