Replace duplicate subtrees

repdupe.go

package main

import (
	"fmt"
)

type Node struct {
	Data  string
	Left  *Node
	Right *Node
}

func (n *Node) Equals(o *Node) bool {
	if n == nil || o == nil {
		return n == o
	}
	return n.Data == o.Data && n.Left == o.Left && n.Right == o.Right
}

type NodeEntry struct {
	Node   *Node
	Parent *Node
}

func uniqueCount(n *Node) int {
	set := make(map[*Node]struct{})
	var recurse func(n *Node)
	recurse = func(n *Node) {
		if n == nil {
			return
		}
		set[n] = struct{}{}
		recurse(n.Left)
		recurse(n.Right)
	}
	recurse(n)
	return len(set)
}

func buildLevelSets(n *Node, s *[][]NodeEntry, level int, parent *Node) {
	if n == nil {
		return
	}
	if len(*s) <= level {
		*s = append(*s, nil)
	}
	(*s)[level] = append((*s)[level], NodeEntry{
		Node:   n,
		Parent: parent,
	})
	buildLevelSets(n.Left, s, level+1, n)
	buildLevelSets(n.Right, s, level+1, n)
}

func main() {

	n := &Node{
		Data:  "A",
		Left:  &Node{Data: "B", Left: &Node{Data: "C"}},
		Right: &Node{Data: "B", Left: &Node{Data: "C"}},
	}

	fmt.Println("Unodes:", uniqueCount(n))
	var levels [][]NodeEntry
	buildLevelSets(n, &levels, 0, nil)

	collapseLevelSets(levels)

	fmt.Println("Unodes:", uniqueCount(n))
}

func collapseLevelSets(e [][]NodeEntry) {
	for i := len(e) - 1; i >= 0; i-- {
		collapseLevel(e[i])
	}
}

func collapseLevel(e []NodeEntry) {
	lookup := make(map[string]*Node)

	for _, ne := range e {
		existing, ok := lookup[ne.Node.Data]
		if !ok {
			lookup[ne.Node.Data] = ne.Node
			continue
		}

		if existing.Equals(ne.Node) {
			replaceChild(ne.Parent, existing)
		}
	}
}

func replaceChild(parent *Node, new *Node) {
	if new.Equals(parent.Left) {
		parent.Left = new
	}
	if new.Equals(parent.Right) {
		parent.Right = new
	}
}