Utility for performing a random split for text data that is already in
    bag-of-words format. For each (word, count) pair in a particular element,
    the counts are uniformly partitioned in either a training set or a test
    set.

    Parameters
    ----------
    dataset : SArray of type dict, SFrame with columns of type dict
        A data set in bag-of-words format.

    prob : float, optional
        Probability for sampling a word to be placed in the test set.

    Returns
    -------
    train, test : SArray
        Two data sets in bag-of-words format, where the combined counts are
        equal to the counts in the original data set.

    Examples
    --------
    >>> docs = turicreate.SArray([{'are':5, 'you':3, 'not': 1, 'entertained':10}])
    >>> train, test = turicreate.text_analytics.random_split(docs)
    >>> print(train)
    [{'not': 1.0, 'you': 3.0, 'are': 3.0, 'entertained': 7.0}]
    >>> print(test)
    [{'are': 2.0, 'entertained': 3.0}]

def random_split(dataset, prob=.5):
    """
    Utility for performing a random split for text data that is already in
    bag-of-words format. For each (word, count) pair in a particular element,
    the counts are uniformly partitioned in either a training set or a test
    set.

    Parameters
    ----------
    dataset : SArray of type dict, SFrame with columns of type dict
        A data set in bag-of-words format.

    prob : float, optional
        Probability for sampling a word to be placed in the test set.

    Returns
    -------
    train, test : SArray
        Two data sets in bag-of-words format, where the combined counts are
        equal to the counts in the original data set.

    Examples
    --------
    >>> docs = turicreate.SArray([{'are':5, 'you':3, 'not': 1, 'entertained':10}])
    >>> train, test = turicreate.text_analytics.random_split(docs)
    >>> print(train)
    [{'not': 1.0, 'you': 3.0, 'are': 3.0, 'entertained': 7.0}]
    >>> print(test)
    [{'are': 2.0, 'entertained': 3.0}]
    """

    def grab_values(x, train=True):
        if train:
            ix = 0
        else:
            ix = 1
        return dict([(key, value[ix]) for key, value in six.iteritems(x) \
                    if value[ix] != 0])

    def word_count_split(n, p):
        num_in_test = 0
        for i in range(n):
            if random.random() < p:
                num_in_test += 1
        return [n - num_in_test, num_in_test]

    # Get an SArray where each word has a 2 element list containing
    # the count that will be for the training set and the count that will
    # be assigned to the test set.
    data = dataset.apply(lambda x: dict([(key, word_count_split(int(value), prob)) \
                                for key, value in six.iteritems(x)]))

    # Materialize the data set
    data.__materialize__()

    # Grab respective counts for each data set
    train = data.apply(lambda x: grab_values(x, train=True))
    test  = data.apply(lambda x: grab_values(x, train=False))

    return train, test