word-wrap.go

type TrieNode struct {
    pool map[rune]*TrieNode
}

func NewTrieNode() *TrieNode {
    return &TrieNode{pool: make(map[rune]*TrieNode)}
}

func (t *TrieNode) Insert(word string) {
    ptr := t
    for _, ch := range word {
        if _, ok := ptr.pool[ch]; !ok {
            ptr.pool[ch] = NewTrieNode()
        }
        ptr = ptr.pool[ch]
    }
    ptr.pool[0] = nil
}

type QueueItem struct {
    ptr *TrieNode
    offset int
    acc []int
}

func wordBreak(s string, wordDict []string) []string {
    t := NewTrieNode()
    for _, word := range wordDict {
        t.Insert(word)
    }
    ix := [][]int{}
    q := make([]QueueItem, 0, len(s))
    q = append(q, QueueItem{t, 0, []int{}})
    var head QueueItem
    for len(q) > 0 {
        head, q = q[0], q[1:]
        if head.offset == len(s) {
            if _, ok := head.ptr.pool[0]; ok {
                ix = append(ix, head.acc)
            }
            continue
        }
        if _, ok := head.ptr.pool[0]; ok {
            q = append(q, QueueItem{
                ptr: t,
                offset: head.offset,
                acc: append(head.acc, head.offset),
            })
        }
        if next, ok := head.ptr.pool[rune(s[head.offset])]; ok {
            q = append(q, QueueItem{
                ptr: next,
                offset: head.offset + 1,
                acc: head.acc,
            })
        }
    }
    
    res := make([]string, 0, len(ix))
    
    for _, ixset := range ix {
        b := make([]string, 0, len(ixset)+1)
        prev := 0
        for _, off := range ixset {
            b = append(b, s[prev:off])
            prev = off
        }
        b = append(b, s[prev:len(s)])
        res = append(res, strings.Join(b, " "))
    }
    
    return res
}