cses

My solution for the coding challenges of cses.fi
git clone https://git.tronto.net/cses
Download | Log | Files | Refs | README

word_combinations_1731.cpp (1243B)


      1 #include <iostream>
      2 #include <queue>
      3 #include <string>
      4 #include <string_view>
      5 #include <vector>
      6 
      7 struct TrieNode {
      8 	char c;
      9 	size_t d;
     10 	bool isend;
     11 	std::vector<TrieNode> next;
     12 
     13 	TrieNode(char x, size_t y) : c{x}, d{y}, isend{false}, next() {}
     14 };
     15 
     16 void push(std::string_view s, TrieNode& t) {
     17 	if (s.empty()) {
     18 		t.isend = true;
     19 		return;
     20 	}
     21 	for (auto& u : t.next) {
     22 		if (s[0] == u.c) {
     23 			push(s.substr(1), u);
     24 			return;
     25 		}
     26 	}
     27 	t.next.push_back(TrieNode(s[0], t.d+1));
     28 	push(s.substr(1), t.next.back());
     29 }
     30 
     31 const TrieNode* next(const TrieNode* w, char c) {
     32 	for (size_t i = 0; i < w->next.size(); i++)
     33 		if (w->next[i].c == c)
     34 			return &w->next[i];
     35 	return nullptr;
     36 }
     37 
     38 int f(std::vector<int>& t, std::string_view s, const TrieNode& d, size_t i) {
     39 	static constexpr long long mod = 1e9+7;
     40 
     41 	if (t[i] != -1) return t[i];
     42 
     43 	t[i] = 0;
     44 	for (const TrieNode* w = &d; w != nullptr; w = next(w, s[w->d]))
     45 		if (w->isend)
     46 			t[i] = (t[i] + f(t, s.substr(w->d), d, i+w->d)) % mod;
     47 	return t[i];
     48 }
     49 
     50 int main() {
     51 	std::string s, u;
     52 	size_t k;
     53 	std::cin >> s >> k;
     54 	TrieNode d('\0', 0);
     55 	for (size_t i = 0; i < k; i++) {
     56 		std::cin >> u;
     57 		push(u, d);
     58 	}
     59 
     60 	std::vector<int> t(s.size()+1, -1);
     61 	t[s.size()] = 1;
     62 	std::cout << f(t, s, d, 0) << "\n";
     63 }