RosettaCodeData/Task/LZW-compression/Go/lzw-compression.go

96 lines
2.1 KiB
Go

package main
import (
"fmt"
"log"
"strings"
)
// Compress a string to a list of output symbols.
func compress(uncompressed string) []int {
// Build the dictionary.
dictSize := 256
// We actually want a map of []byte -> int but
// slices are not acceptable map key types.
dictionary := make(map[string]int, dictSize)
for i := 0; i < dictSize; i++ {
// Ugly mess to work around not having a []byte key type.
// Using `string(i)` would do utf8 encoding for i>127.
dictionary[string([]byte{byte(i)})] = i
}
var result []int
var w []byte
for i := 0; i < len(uncompressed); i++ {
c := uncompressed[i]
wc := append(w, c)
if _, ok := dictionary[string(wc)]; ok {
w = wc
} else {
result = append(result, dictionary[string(w)])
// Add wc to the dictionary.
dictionary[string(wc)] = dictSize
dictSize++
//w = []byte{c}, but re-using wc
wc[0] = c
w = wc[:1]
}
}
if len(w) > 0 {
// Output the code for w.
result = append(result, dictionary[string(w)])
}
return result
}
type BadSymbolError int
func (e BadSymbolError) Error() string {
return fmt.Sprint("Bad compressed symbol ", int(e))
}
// Decompress a list of output symbols to a string.
func decompress(compressed []int) (string, error) {
// Build the dictionary.
dictSize := 256
dictionary := make(map[int][]byte, dictSize)
for i := 0; i < dictSize; i++ {
dictionary[i] = []byte{byte(i)}
}
var result strings.Builder
var w []byte
for _, k := range compressed {
var entry []byte
if x, ok := dictionary[k]; ok {
//entry = x, but ensuring any append will make a copy
entry = x[:len(x):len(x)]
} else if k == dictSize && len(w) > 0 {
entry = append(w, w[0])
} else {
return result.String(), BadSymbolError(k)
}
result.Write(entry)
if len(w) > 0 {
// Add w+entry[0] to the dictionary.
w = append(w, entry[0])
dictionary[dictSize] = w
dictSize++
}
w = entry
}
return result.String(), nil
}
func main() {
compressed := compress("TOBEORNOTTOBEORTOBEORNOT")
fmt.Println(compressed)
decompressed, err := decompress(compressed)
if err != nil {
log.Fatal(err)
}
fmt.Println(decompressed)
}