import collections


def build_dataset(input_words, vocabulary_size):
    """ Build the appropriate datasets for the desired vocabulary size given a corpus of input words """

    # Extract the n-most common words from the input (given a desired vocabulary size)
    vocabulary_words = collections.Counter(input_words).most_common(vocabulary_size)
    # Create a dictionary of indexes for words in the vocabulary, in descending order of most common
    vocab = dict()
    for word, _ in vocabulary_words:
        vocab[word] = len(vocab) + 1
    # Create the input corpus data - the index in the vocabulary of every known word in the input text
    corpus_data = list()
    for word in input_words:
        if word in vocab:
            index = vocab[word]
        else:
            index = 0
        corpus_data.append(index)
    return corpus_data, dict(zip(vocab.values(), vocab.keys()))