diff --git a/huffman_test.go b/huffman_test.go index a5ca063..289b82f 100644 --- a/huffman_test.go +++ b/huffman_test.go @@ -1,6 +1,7 @@ package huffman import ( + "bytes" "fmt" "testing" @@ -49,6 +50,12 @@ func TestBuildTree0(t *testing.T) { assert.Equal(t, int64(2), n.r.l.l.key) assert.Equal(t, int64(3), n.r.l.r.key) assert.Equal(t, int64(6), n.r.r.key) + + w := bytes.NewBufferString("") + printTree(w, n) + if debug { + fmt.Println(w) + } } func TestBuildTree1(t *testing.T) { @@ -68,6 +75,12 @@ func TestBuildTree1(t *testing.T) { assert.Equal(t, int64(58), n.key) assert.Equal(t, int64(25), n.l.key) assert.Equal(t, int64(33), n.r.key) + + w := bytes.NewBufferString("") + printTree(w, n) + if debug { + fmt.Println(w) + } } func TestGenerateTable(t *testing.T) { diff --git a/utils.go b/utils.go new file mode 100644 index 0000000..5d1c48c --- /dev/null +++ b/utils.go @@ -0,0 +1,24 @@ +package huffman + +import ( + "fmt" + "io" +) + +func printNode(w io.Writer, n *Node) { + if n.typ == MID { + fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", n.String(), n.l.String(), n.r.String()) + printNode(w, n.l) + printNode(w, n.r) + } else if n.typ == LEAF { + fmt.Fprintf(w, "\"%v\" [style=filled];\n", n.String()) + } +} +func printTree(w io.Writer, n *Node) { + fmt.Fprintf(w, `digraph hierarchy { + node [fontname=Monospace,fontsize=10,shape=box] + `) + + printNode(w, n) + fmt.Fprintf(w, "}\n") +}