class TrieNode: def __init__(self, val = ''): self.child = [] self.is_word = 0 self.val = val self.t = 1 self.child_map = {} class Trie: def __init__(self): self.root = TrieNode() self.cnt = 0 def insert(self, word: str) -> None: parent_node = self.root for idx, letter in enumerate(word): if letter not in parent_node.child_map: new_node = TrieNode(letter) parent_node.child.append(new_node) parent_node.child_map[letter] = len(parent_node.child) - 1 parent_node = new_node else: idx = parent_node.child_map[letter] parent_node = parent_node.child[idx] if parent_node.is_word != 0: parent_node.t += 1 parent_node.is_word = self.cnt + 1 self.cnt += 1 def search(self, word: str) : parent_node = self.root for idx, letter in enumerate(word): if letter not in parent_node.child_map: return None idx = parent_node.child_map[letter] parent_node = parent_node.child[idx] if parent_node.is_word != 0: return parent_node return None def startsWith(self, prefix: str) -> bool: parent_node = self.root for idx, letter in enumerate(prefix): if letter not in parent_node.child_map: return False idx = parent_node.child_map[letter] parent_node = parent_node.child[idx] return True class Solution: def findSubstring(self, s: str, words: list[str]) -> list[int]: len_word = len(words[0]) tree = Trie() is_word = [] rlt = [] state = [] word_cnt = {} if len_word > len(s): return [] for word in words: tree.insert(word) r = tree.search(word) if r != None: word_cnt[r.is_word] = r.t # print(word_cnt) for idx, ch in enumerate(s): if idx + len_word > len(s): break waiting = s[idx:idx+len_word] r = tree.search(waiting) if r != None: is_word.append(r.is_word) else: is_word.append(0) for i in range(len_word - 1) : is_word.append(0) for i in range(len_word): if is_word[i] != 0: tmp_dic = {is_word[i]: [i]} state.append((tmp_dic,i,i)) else: state.append(({},i,i)) if len(words) == 1 and words[0] == s: return [0] # is_word[idx] = value words index / 0 if it is not a word # state[dic_idx][0] for i in range(len(is_word[len_word:])): idx = i + len_word dic_idx = idx % len_word # print(f"current index: {idx}, {idx % len_word}") # print(f"current dic: {state[dic_idx]}") # print(f"is_word[idx]: {is_word[idx]}") if is_word[idx] in state[dic_idx][0] and len(state[dic_idx][0][is_word[idx]]) >= word_cnt[is_word[idx]]: # print("condition 1") tmp = state[dic_idx][0][is_word[idx]][0] tmp_l = [] for key in state[dic_idx][0]: for i, v in enumerate(state[dic_idx][0][key]): if v <= tmp: tmp_l.append((key, i)) for (key, i) in tmp_l: state[dic_idx][0][key].pop(i) if len(state[dic_idx][0][key]) == 0: del state[dic_idx][0][key] if is_word[idx] not in state[dic_idx][0]: state[dic_idx][0][is_word[idx]] = [idx] else: state[dic_idx][0][is_word[idx]].append(idx) state[dic_idx] = (state[dic_idx][0], tmp + len_word, state[dic_idx][2]) elif is_word[idx] != 0: # print("condition 2") if len(state[dic_idx][0]) == 0: state[dic_idx] = ({}, idx, idx) if is_word[idx] in state[dic_idx][0]: state[dic_idx][0][is_word[idx]].append(idx) else: state[dic_idx][0][is_word[idx]]=[idx] elif is_word[idx] == 0: # print("condition 3") state[dic_idx] = ({}, idx, idx) state[dic_idx] = (state[dic_idx][0],state[dic_idx][1], idx) s = 0 for k in state[dic_idx][0]: s += len(state[dic_idx][0][k]) # if len(state[dic_idx][0]) == len(words): if s == len(words): # print(f"win, {idx}") rlt.append(state[dic_idx][1]) del state[dic_idx][0][is_word[state[dic_idx][1]]] state[dic_idx] = (state[dic_idx][0], state[dic_idx][1] + len_word, state[dic_idx][2]) # print(f"after fixed: {state[dic_idx]}") # print() return rlt sol = Solution() print(sol.findSubstring("barfoothefoobarman", words=["foo", "bar"])) print(sol.findSubstring("barfoofoobarthefoobarman", words=["foo", "bar", "the"])) print(sol.findSubstring("barbarbbar", words=["arb", "bar"])) print(sol.findSubstring("wordgoodgoodgoodbestword", words=["word", "good", "good", "best"])) print(sol.findSubstring("aaa", words=["a"])) print(sol.findSubstring("word", words=["word"]))