CTCI_3.3..go

package main

type stack struct {
	data []interface{}
}

func newStack() *stack {
	return &stack{
		data: make([]interface{}, 0),
	}
}

func (s *stack) push(d interface{}) {
	s.data = append(s.data, d)
}

func (s *stack) pop() interface{} {
	last := s.data[len(s.data)-1]
	s.data = s.data[:len(s.data)-1]

	return last
}

type setOfStacks struct {
	stacks                []*stack
	numOfStack, threshold int
}

func newSetOfStacks(th int) *setOfStacks {
	s := newStack()
	stacks := make([]*stack, 0)
	stacks = append(stacks, s)

	return &setOfStacks{
		stacks:     stacks,
		threshold:  th,
		numOfStack: 1,
	}
}

func (s *setOfStacks) push(d interface{}) {
	if len(s.stacks) == 0 || len(s.stacks[s.numOfStack-1].data) == s.threshold {
		lastStack := newStack()
		lastStack.push(d)
		s.stacks = append(s.stacks, lastStack)
		s.numOfStack++
		return
	}

	lastStack := s.stacks[s.numOfStack-1]
	if len(lastStack.data) < s.threshold {
		lastStack.push(d)
	}
}

func (s *setOfStacks) pop() interface{} {
	lastStack := s.stacks[s.numOfStack-1]
	lastData := lastStack.pop()

	if len(lastStack.data) == 0 {
		s.stacks = s.stacks[:len(s.stacks)-1]
		s.numOfStack--
	}

	return lastData
}

func main() {
	ss := newSetOfStacks(3)

	ss.push(1)
	ss.push(2)
	ss.push(3)
	ss.push(4)

	println("len of stacks: ", len(ss.stacks[ss.numOfStack-1].data))
	println("num of stacks: ", ss.numOfStack)

	println(ss.pop().(int))

	println("len of stacks: ", len(ss.stacks[ss.numOfStack-1].data))
	println("num of stacks: ", ss.numOfStack)
}