CSES Solutions – Word Combinations

Last Updated : 02 Apr, 2024

You are given a string S of length N and a dictionary arr[] containing K words. In how many ways can you create the string using the words?

Examples:

Input: S = “ababc”, K = 4, arr[] = {ab, abab, c, cb}
Output: 2
Explanation: The possible ways are “ab” + “ab” + “c” and “abab” + “c”.

Input: S = “geeksforgeeks”, K = 3, arr[] = {geeks, for, geeks}
Output: 2
Explanation: The possible ways are “geeks” + “for” + “geeks” and “geeks” + “for” + “geeks” .

Approach: To solve the problem, follow the below idea:

The problem can be solved by Dynamic Programming and Tries. We can maintain a dp[] array such that dp[i] stores the number of ways to form the substring S[i…N-1]. We can travverse the string S from right to left and for every index i, dp[i] is the sum of all dp[j + 1] such that S[i…j] is present in the dictionary. In order to check whether dp[i…j] is present in the dictionary or not, we can insert all the words in a trie and mark the ending of each word with a flag. So, for every index i, we iterate j from i to N-1 and if S[i…j] is present in the dictionary, we add dp[j] to dp[i]. After all the iterations, dp[0] will store the answer.

Step-by-Step algorithm:

Declare a dp[] array to store the number of ways to form the string from each index.
Declare the Trie data structure to store all the words in a way that allows efficient search and retrieval.
Insert each word into the Trie: For each word in the given list, insert it into the Trie using the insertWordInTrie() function. This function iterates over each character in the word, and for each character, it checks if a node already exists in the Trie for this character. If not, it creates a new node. It then marks the last node of each word as the end of a word.
For each index i from s.size() – 1 to 0, calculate dp[i] using the countWays() function.
The countWays() function iterates over the string from the current index, and for each character, it checks if a node exists in the Trie for this character.
- If not, it returns the current number of ways.
- If a node does exist and is marked as end of word then it adds the number of ways to form the string from the next index to the current number of ways.
The answer is the number of ways to form the string from index 0, which is stored in dp[0].

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
using namespace std;

#define ll long long
const int mod = 1e9 + 7;

// Declare the DP array and the Trie data structure
vector<int> dp(5005);
vector<vector<int> > trie(1e6 + 5, vector<int>(26));
vector<bool> isEndOfWord(1e6 + 5);
int trieNodeCount = 0;

// Function to insert a word into the Trie
void insertWordInTrie(string& word)
{
    int currentNode = 0;
    for (char& ch : word) {
        if (!trie[currentNode][ch - 'a'])
            trie[currentNode][ch - 'a'] = ++trieNodeCount;
        currentNode = trie[currentNode][ch - 'a'];
    }
    isEndOfWord[currentNode] = true;
}

// Function to count the number of ways to form the string
// from index 'start'
int countWays(int start, string& S)
{
    int currentNode = 0, ways = 0;
    for (int i = start; i < S.size(); i++) {
        if (!trie[currentNode][S[i] - 'a'])
            return ways;
        currentNode = trie[currentNode][S[i] - 'a'];
        // If a word ends here, add the number of ways from the next index
        if (isEndOfWord[currentNode])
            ways = (ways + dp[i + 1]) % mod; 
    }
    return ways;
}

void solve(string& S, int K, string word[])
{
    for (int i = 0; i < K; i++) {
        // Insert each word into the Trie
        insertWordInTrie(word[i]); 
    }
    // Base case: One way to form an empty string
    dp[S.size()] = 1; 
    for (int i = S.size() - 1; i >= 0; i--) {
        // Fill the DP table from right to left
        dp[i] = countWays(i, S);
    }
    // The answer is the number of ways to form the string from index 0
    cout << dp[0]; 
}
// Driver Code
int main()
{
    // Input from user
    string S = "ababc";
    int K = 4;
    string arr[] = { "ab", "abab", "c", "cb" };
    solve(S, K, arr);
    return 0;
}

Java

import java.util.Arrays;

public class Main {
    static final int MOD = (int) 1e9 + 7;
    static int[] dp = new int[5005];
    static int[][] trie = new int[(int) 1e6 + 5][26];
    static boolean[] isEndOfWord = new boolean[(int) 1e6 + 5];
    static int trieNodeCount = 0;

    // Function to insert a word into the Trie
    static void insertWordInTrie(String word) {
        int currentNode = 0;
        for (char ch : word.toCharArray()) {
            if (trie[currentNode][ch - 'a'] == 0)
                trie[currentNode][ch - 'a'] = ++trieNodeCount;
            currentNode = trie[currentNode][ch - 'a'];
        }
        isEndOfWord[currentNode] = true;
    }

    // Function to count the number of ways to form the string
    // from index 'start'
    static int countWays(int start, String S) {
        int currentNode = 0, ways = 0;
        for (int i = start; i < S.length(); i++) {
            if (trie[currentNode][S.charAt(i) - 'a'] == 0)
                return ways;
            currentNode = trie[currentNode][S.charAt(i) - 'a'];
            // If a word ends here, add the number of ways from the next index
            if (isEndOfWord[currentNode])
                ways = (ways + dp[i + 1]) % MOD;
        }
        return ways;
    }

    static void solve(String S, int K, String[] words) {
        for (int i = 0; i < K; i++) {
            // Insert each word into the Trie
            insertWordInTrie(words[i]);
        }
        // Base case: One way to form an empty string
        dp[S.length()] = 1;
        for (int i = S.length() - 1; i >= 0; i--) {
            // Fill the DP table from right to left
            dp[i] = countWays(i, S);
        }
        // The answer is the number of ways to form the string from index 0
        System.out.println(dp[0]);
    }

    // Driver Code
    public static void main(String[] args) {
        // Input from user
        String S = "ababc";
        int K = 4;
        String[] arr = { "ab", "abab", "c", "cb" };
        solve(S, K, arr);
    }
}

// This code is contributed by rambabuguphka

Python

MOD = 10**9 + 7

# Declare the DP array and the Trie data structure
dp = [0] * 5005
trie = [[0] * 26 for _ in range(10**6 + 5)]
isEndOfWord = [False] * (10**6 + 5)
trieNodeCount = 0

# Function to insert a word into the Trie
def insertWordInTrie(word):
    global trieNodeCount
    currentNode = 0
    for ch in word:
        index = ord(ch) - ord('a')
        if not trie[currentNode][index]:
            trieNodeCount += 1
            trie[currentNode][index] = trieNodeCount
        currentNode = trie[currentNode][index]
    isEndOfWord[currentNode] = True

# Function to count the number of ways to form the string
# from index 'start'
def countWays(start, S):
    currentNode = 0
    ways = 0
    for i in range(start, len(S)):
        index = ord(S[i]) - ord('a')
        if not trie[currentNode][index]:
            return ways
        currentNode = trie[currentNode][index]
        # If a word ends here, add the number of ways from the next index
        if isEndOfWord[currentNode]:
            ways = (ways + dp[i + 1]) % MOD
    return ways

def solve(S, K, word):
    for w in word:
        # Insert each word into the Trie
        insertWordInTrie(w)
    # Base case: One way to form an empty string
    dp[len(S)] = 1
    for i in range(len(S) - 1, -1, -1):
        # Fill the DP table from right to left
        dp[i] = countWays(i, S)
    # The answer is the number of ways to form the string from index 0
    print(dp[0])

# Driver Code
if __name__ == "__main__":
    # Input from user
    S = "ababc"
    K = 4
    arr = ["ab", "abab", "c", "cb"]
    solve(S, K, arr)

using System;

public class Gfg
{
    const int MOD = 1000000007;
    static int[] dp = new int[5005];
    static int[,] trie = new int[1000005, 26];
    static bool[] isEndOfWord = new bool[1000005];
    static int trieNodeCount = 0;

    // Function to insert a word into the Trie
    static void InsertWordInTrie(string word)
    {
        int currentNode = 0;
        foreach (char ch in word)
        {
            if (trie[currentNode, ch - 'a'] == 0)
                trie[currentNode, ch - 'a'] = ++trieNodeCount;
            currentNode = trie[currentNode, ch - 'a'];
        }
        isEndOfWord[currentNode] = true;
    }

    // Function to count the number of ways to form the string
    // from index 'start'
    static int CountWays(int start, string S)
    {
        int currentNode = 0, ways = 0;
        for (int i = start; i < S.Length; i++)
        {
            if (trie[currentNode, S[i] - 'a'] == 0)
                return ways;
            currentNode = trie[currentNode, S[i] - 'a'];
            // If a word ends here, add the number of ways from the next index
            if (isEndOfWord[currentNode])
                ways = (ways + dp[i + 1]) % MOD;
        }
        return ways;
    }

    static void Solve(string S, int K, string[] words)
    {
        for (int i = 0; i < K; i++)
        {
            // Insert each word into the Trie
            InsertWordInTrie(words[i]);
        }
        // Base case: One way to form an empty string
        dp[S.Length] = 1;
        for (int i = S.Length - 1; i >= 0; i--)
        {
            // Fill the DP table from right to left
            dp[i] = CountWays(i, S);
        }
        // The answer is the number of ways to form the string from index 0
        Console.WriteLine(dp[0]);
    }

    // Driver Code
    public static void Main(string[] args)
    {
        // Input from user
        string S = "ababc";
        int K = 4;
        string[] arr = { "ab", "abab", "c", "cb" };
        Solve(S, K, arr);
    }
}

JavaScript

const mod = 1e9 + 7;

// Declare the DP array and the Trie data structure
let dp = new Array(5005).fill(0);
let trie = new Array(1e6 + 5).fill(null).map(() => new Array(26).fill(0));
let isEndOfWord = new Array(1e6 + 5).fill(false);
let trieNodeCount = 0;

// Function to insert a word into the Trie
function insertWordInTrie(word) {
    let currentNode = 0;
    for (let ch of word) {
        if (!trie[currentNode][ch.charCodeAt(0) - 'a'.charCodeAt(0)]) {
            trie[currentNode][ch.charCodeAt(0) - 'a'.charCodeAt(0)] = ++trieNodeCount;
        }
        currentNode = trie[currentNode][ch.charCodeAt(0) - 'a'.charCodeAt(0)];
    }
    isEndOfWord[currentNode] = true;
}

// Function to count the number of ways to form the string
// from index 'start'
function countWays(start, S) {
    let currentNode = 0;
    let ways = 0;
    for (let i = start; i < S.length; i++) {
        if (!trie[currentNode][S.charCodeAt(i) - 'a'.charCodeAt(0)]) {
            return ways;
        }
        currentNode = trie[currentNode][S.charCodeAt(i) - 'a'.charCodeAt(0)];
        // If a word ends here, add the number of ways from the next index
        if (isEndOfWord[currentNode]) {
            ways = (ways + dp[i + 1]) % mod;
        }
    }
    return ways;
}

function solve(S, K, word) {
    for (let i = 0; i < K; i++) {
        // Insert each word into the Trie
        insertWordInTrie(word[i]);
    }
    // Base case: One way to form an empty string
    dp[S.length] = 1;
    for (let i = S.length - 1; i >= 0; i--) {
        // Fill the DP table from right to left
        dp[i] = countWays(i, S);
    }
    // The answer is the number of ways to form the string from index 0
    console.log(dp[0]);
}

// Driver Code
function main() {
    // Input from user
    let S = "ababc";
    let K = 4;
    let arr = ["ab", "abab", "c", "cb"];
    solve(S, K, arr);
}

main();

// This code is contributed by Ayush Mishra

Output

Time complexity: O(N * N + N * K), where N is the number of length of S and K is the number of words in the dictionary.
Auxiliary Space: O(N * K)