lc616.py

class Solution:
    def build_trie(self, words):
        root = {}
        count = 0
        for word in words:
            parent = root
            for ch in word:
                count += 1
                if ch in parent:
                    parent = parent[ch]
                else:
                    parent[ch] = {}
                    parent = parent[ch]
            else:
                parent["$"] = count
            count = 0
        return root

    def in_trie(self, trie, word):
        parent = trie
        for ch in word:
            if ch not in parent:
                return parent.get('$', 0)
            else:
                parent = parent[ch]
        else:
            return parent.get('$', 0)

    def addBoldTag(self, s, dic):
        cache = [0] * len(s)
        result = []
        trie = self.build_trie(dic)

        for idx in range(len(s)):
            count = self.in_trie(trie, s[idx:])
            for c in range(count):
                cache[idx+c] = 1

        status = 0

        for idx, flag in enumerate(cache):
            if flag and not status:
                result.append("<b>")
                result.append(s[idx])
                status = 1
            elif not flag and status:
                result.append("</b>")
                result.append(s[idx])
                status = 0
            else:
                result.append(s[idx])
        else:
            if status:
                result.append("</b>")

        return "".join(result)