LeetCode 943 "Find the Shortest Superstring"

943.go

func shortestSuperstring(words []string) string {
	n := uint16(len(words))
	lt := makeLenTable(words)
	s := searcher{uint16(n), lt, make([]uint8, n<<n)}

	min := uint8(255)
	first := uint16(0)
	for i := uint16(0); i < n; i++ {
		val := s.search(key(i << n))
		if min > val {
			min = val
			first = i
		}
	}

	seq := s.reconstruct(first)
	// Concatenate
	seqString := make([]string, n)
	for i, w := range seq {
		if i == len(seq)-1 {
			seqString[i] = words[w]
		} else {
			seqString[i] = words[w][:lt.prefix(w, seq[i+1])]
		}
	}
	return strings.Join(seqString, "")
}

type lenTable struct {
	n uint16
	l []uint8
	p []uint8
}

func makeLenTable(words []string) lenTable {
	n := len(words)
	l := make([]uint8, n)
	p := make([]uint8, n*n)
	for i, word1 := range words {
		l[i] = uint8(len(word1))
		for j, word2 := range words {
			for k := 0; k <= len(word1); k++ {
				if strings.HasPrefix(word2, word1[k:]) {
					p[i*n+j] = uint8(k)
					break
				}
			}
		}
	}
	return lenTable{uint16(n), l, p}
}

func (lt lenTable) len(i uint16) uint8 { return lt.l[i] }

// The minimal k such that words[i][k:] is a prefix of words[j].
func (lt lenTable) prefix(i, j uint16) uint8 { return lt.p[i*lt.n+j] }

type key uint16

func encode(n, first uint16, used set) key { return key(first<<n | uint16(used)) }

func (k key) decode(n uint16) (first uint16, used set) { return uint16(k) >> n, set(k) & full(n) }

type set uint16

func full(n uint16) set { return set(1<<n - 1) }

func (s set) has(i uint16) bool { return s&(1<<i) != 0 }

func (s set) with(i uint16) set { return s | (1 << i) }

type searcher struct {
	n  uint16
	lt lenTable
	// key is (next << n) | bitset, where bitset is the set of words *not* to use
	answer []uint8
}

func (s *searcher) search(k key) uint8 {
	if s.answer[k] != 0 {
		return s.answer[k]
	}
	first, used := k.decode(s.n)
	used = used.with(first)
	if used == full(s.n) {
		// No other words to use, just first
		l := s.lt.len(first)
		s.answer[k] = l
		return l
	}
	min := uint8(255)
	for i := uint16(0); i < s.n; i++ {
		if used.has(i) {
			// Already used
			continue
		}
		val := s.lt.prefix(first, i) + s.search(encode(s.n, i, used))
		if min > val {
			min = val
		}
	}
	s.answer[k] = min
	return min
}

func (s *searcher) reconstruct(first uint16) []uint16 {
	seq := make([]uint16, s.n)
	seq[0] = first
	k := encode(s.n, first, 0)

	for d := uint16(1); d < s.n; d++ {
		// The following mirrors search above
		first, used := k.decode(s.n)
		used = used.with(first)
		for i := uint16(0); i < s.n; i++ {
			if used.has(i) {
				continue
			}
			if s.answer[k] == s.lt.prefix(first, i)+s.answer[encode(s.n, i, used)] {
				seq[d] = i
				k = encode(s.n, i, used)
			}
		}
	}
	return seq
}