import random


def find_trigram(vocab, pair=False):
    """Find random set of related words in the vocab dictionary.

    If a valid pair is passed, return the third word in the trigram.
    Otherwise, randomly choose and return a key pair from the dictionary.
    """
    if pair in vocab:
        return random.choice(vocab[pair])
    return random.choice(list(vocab))