Solution to Leetcode problem 1078: Occurrences After Bigram

In this problem, given a sequence of words, you should return all occurences of words that appear after the words first second , where first and second are inputs.

For instance, given the input text: The only thing I know is that I know nothing, and first: I and second: know , the output should be [is, nothing] .

In this problem, we mostly need to track a state given the 2 previous words we have seen, and we need to make sure that each time the 2 previous words are first and second, then we add the next word.

Some tricky cases are if first and second are the same, or if the word added is also a word in the list. In that case we need to be careful keeping the right state.

text: "the the problem is the the the words are repeated"
first: "the"
second: "the"
expected output: ["problem", "the", "words"]

The way I ensure it is quite simple, by going case after case. I do the check in a particular order so that I don’t erase a state that I then want to use. Another possible ways is to use variables such as oldState, and newState, which enforces that the old state is not replaced while you still need it.

Here is the code:

func findOcurrences(
        text string, first string, second string) []string {
    arr := strings.Split(text, " ")
    res := []string{} 
    firstCheck := false
    secondCheck := false
    for _, word := range arr {
        if secondCheck {
            res = append(res, word)
            secondCheck = false
        }
        if firstCheck && strings.Compare(second, word) == 0 {
            secondCheck = true
            firstCheck = false
        } else if firstCheck {
            firstCheck = false
        }
        if strings.Compare(first, word) == 0 {
            firstCheck = true
        }
    }
    return res
}