Occurrences After Bigram

Problem

Given two strings first and second, consider occurrences in some text of the form "first second third", where second comes immediately after first, and third comes immediately after second.

Return an array of all the words third for each occurrence of "first second third".

Examples

Example 1

Input: text = "alice is a good girl she is a good student", first = "a", second = "good"
Output: ["girl","student"]

Example 2

Input: text = "we will we will rock you", first = "we", second = "will"
Output: ["we","rock"]

Constraints

• 1 <= text.length <= 1000
• text consists of lowercase English letters and spaces.
• All the words in text are separated by a single space.
• 1 <= first.length, second.length <= 10
• first and second consist of lowercase English letters.
• text will not have any leading or trailing spaces.

Solution

Method 1 – Sliding Window

Intuition

We want to find all words that come immediately after a given bigram (pair of words). By sliding a window of size 3 over the words, we can check if the first two match the bigram and collect the third.

Approach

1. Split the text into words.
2. Iterate through the words with a window of size 3.
3. If the first and second word in the window match first and second, add the third word to the result.
4. Return the result list.

Code

C++

class Solution {
public:
    vector<string> findOcurrences(string text, string first, string second) {
        vector<string> ans;
        istringstream iss(text);
        vector<string> words{istream_iterator<string>{iss}, {}};
        for (int i = 2; i < words.size(); ++i) {
            if (words[i-2] == first && words[i-1] == second)
                ans.push_back(words[i]);
        }
        return ans;
    }
};

Go

func findOcurrences(text, first, second string) []string {
    words := strings.Fields(text)
    ans := []string{}
    for i := 2; i < len(words); i++ {
        if words[i-2] == first && words[i-1] == second {
            ans = append(ans, words[i])
        }
    }
    return ans
}

Java

class Solution {
    public String[] findOcurrences(String text, String first, String second) {
        String[] words = text.split(" ");
        List<String> ans = new ArrayList<>();
        for (int i = 2; i < words.length; i++) {
            if (words[i-2].equals(first) && words[i-1].equals(second))
                ans.add(words[i]);
        }
        return ans.toArray(new String[0]);
    }
}

Kotlin

class Solution {
    fun findOcurrences(text: String, first: String, second: String): Array<String> {
        val words = text.split(" ")
        val ans = mutableListOf<String>()
        for (i in 2 until words.size) {
            if (words[i-2] == first && words[i-1] == second)
                ans.add(words[i])
        }
        return ans.toTypedArray()
    }
}

Python

class Solution:
    def findOcurrences(self, text: str, first: str, second: str) -> list[str]:
        words = text.split()
        ans = []
        for i in range(2, len(words)):
            if words[i-2] == first and words[i-1] == second:
                ans.append(words[i])
        return ans

Rust

impl Solution {
    pub fn find_ocurrences(text: String, first: String, second: String) -> Vec<String> {
        let words: Vec<&str> = text.split_whitespace().collect();
        let mut ans = Vec::new();
        for i in 2..words.len() {
            if words[i-2] == first && words[i-1] == second {
                ans.push(words[i].to_string());
            }
        }
        ans
    }
}

TypeScript

class Solution {
    findOcurrences(text: string, first: string, second: string): string[] {
        const words = text.split(' ')
        const ans: string[] = []
        for (let i = 2; i < words.length; i++) {
            if (words[i-2] === first && words[i-1] === second) ans.push(words[i])
        }
        return ans
    }
}

Complexity

• ⏰ Time complexity: O(n), where n is the number of words in text. Each word is visited once.
• 🧺 Space complexity: O(n), for storing the words and the result.