[자료구조][Go] Priority Queue

코드 작성

├── collections
│   ├── heap     		<--
│   ├── segementTree
│   └── stack
└── pkg
    └── mod
        └── cache

heap

우선순위큐는 힙이라는 자료구조를 이용한다. 힙을 사용할 경우 요소를 추가,삭제하는 연산을 O(log2 N)으로 처리할 수 있다.

힙은 완전이진트리이기 때문에 배열을 사용해서 구현할 수 있다. 이때 인덱스 0은 자식노드 계산을 쉽게하기 위해서 사용되지 않는다. 자식노드를 계산하는 법은 부모노드의 인덱스를 i라고 했을때 2i , 2i+1이 각각 왼쪽,오른쪽 자식이 된다.

Testify

이번에는 Testify라는 툴을 이용해서 테스트를 진행해볼 예정이다. Testify는 “assertion, mocking, test suite” 기능을 제공한다. 이것을 통해 조금 더 효율적으로 테스트를 진행할수 있다.

new 구현

heap_test.go

func TestNewHeap(t *testing.T) {
	h := NewHeap()
	fmt.Println(reflect.TypeOf(h))
	assert.IsType(t, &heap{}, h)
}

heap.go

type heap struct {
	node []int
}

func NewHeap() *heap {
	return &heap{[]int{0}}
}

인덱스 1번부터 사용하기 때문에 인덱스0번에 숫자하나를 채운다.그리고 타입검사를 해서 타입이 같은지 확인한다.

heap_test.go

type TestSuite struct {
	suite.Suite
	h *heap
}

func TestRunSuit(t *testing.T) {
	suite.Run(t, new(TestSuite))
}
// 테스트가 실행될때마다 NewHeap 실행
func (ts *TestSuite) SetupTest() {
	ts.h = NewHeap()
}

테스트를 통과했다면 이제 Testsuite를 만들어 테스트가 실행될때마다 새로운 heap을 생성하도록 한다.

add 구현

heap_test.go

func (ts *TestSuite) TestHeap_Add() {
	ts.h.Add(1)
	ts.h.Add(2)
	ts.h.Add(3)
	ts.Equal(len(ts.h.node)-1, 3)
	assert.Equal(ts.T(), len(ts.h.node)-1, 3)
}

Add를 했을때 값이 추가됐는지 확인하는 간단한 테스트 코드를 일단 작성하고 통과하면 일단 Add는 거기서 끝내고 필요할 때 내용을 더 추가한다.

heap.go

// Add heap에 요소를 추가한다.
// num: 추가하고 싶은 값
func (s *heap) Add(num int) {
	// 맨 마지막 인덱스에 요소 추가
	s.node = append(s.node, num)
}

poll 구현

heap_test.go

func (ts *TestSuite) TestHeap_Poll() {
	ts.h.Add(1)
	ts.h.Add(3)
	ts.h.Add(10)
	result := ts.h.Poll()

	assert.Equal(ts.T(), result, 10)
}

poll은 heap에 있는 값 중 최댓값(혹은 최솟값)을 반환하고 그 값을 제거해야 한다. 내부적으로 heap에 있는 노드들의 인덱스를 조정해 주어야 한다.

구현해야할 연산은 두가지 이다.

• 요소를 추가했을 때 크기를 비교해서 자식노드와 부모노드를 비교해서 조건에 맞으면(부모노드<자식노드) 부모노드와 자식노드위치를 바꾼다.
• 요소를 삭제했을 때 맨마지막 요소를 루트노드에 올린 후 자식노드와 비교해서 조건에 맞으면 부모노드와 자식노드 위치를 바꾼다.

Add()부터 다시 정리해보면

1. heap배열 마지막에 값을 추가
2. 부모노드와 비교해서 조건에 맞으면 자식노드와 교환
3. 루트노드까지 반복

heap.go

// Add heap에 요소를 추가한다.
// num: 추가하고 싶은 값
func (s *heap) Add(num int) {
	// 맨 마지막 인덱스에 요소 추가
	s.node = append(s.node, num)
	// 연산
	index := len(s.node) - 1
	for index > 1 {
		parent := index / 2
		// 부모보다 크면
		if s.node[parent] < s.node[index] {
        	// 스왑
			temp := s.node[parent]
			s.node[parent] = s.node[index]
			s.node[index] = temp
		}
		index = parent
	}
}

다음 Poll()을 구현하는 순서는 다음과 같다.

1. 맨마지막 노드를 루트노드로 옮기고 heap 크기를 하나 줄인다.
2. 왼쪽 자식노드부터 비교한다.오른쪽 자식 노드와 왼쪽 자식 노드 중 더 큰 노드의 값과 교환한다.
3. len(heap)만큼 반복한다.

heap.go

// Poll 요소를 반환하고 삭제한다.
func (s *heap) Poll() int {
	result := s.node[1]
	//맨마지막 요소 루트노드로
	s.node[1] = s.node[len(s.node)-1]
	// 크기 줄이기
	s.node = append(s.node[0 : len(s.node)-1])
	parent := 1

	for parent*2 < len(s.node) {
		//왼쪽부터
		min := s.node[parent*2]
		minIndex := parent * 2
		//오른쪽이 더크면 변경
		if parent*2+1 < len(s.node) && min < s.node[parent*2+1] {
			min = s.node[parent*2+1]
			minIndex = parent*2 + 1
		}
		//부모가 더크면 종료
		if min < s.node[parent] {
			break
		}
		//스왑
		temp := s.node[minIndex]
		s.node[minIndex] = s.node[parent]
		s.node[parent] = temp
		//밑으로
		parent = minIndex
	}
	return result
}

최소힙, 최대힙 나누기

사실 heap에는 두가지 종류가 있다.바로 최대힙과 최소힙이다.

• 최대힙: 부모노드가 자식노드보다 크거나 같다
• 최소힙: 부모노드가 자식노드보다 작거나 같다

최소힙,최대힙 각각 따로 만들면 되지만 중복 코드가 발생 한다. 그래서 우선순위 큐를 좀 더 유연하게 만들 필요가 있다.

// Add heap에 요소를 추가한다.
// num: 추가하고 싶은 값
func (s *heap) Add(num int) {
	// 맨 마지막 인덱스에 요소 추가
	s.node = append(s.node, num)
	// 연산
	index := len(s.node) - 1
	for index > 1 {
		parent := index / 2

		if '함수를 만들어서 관리할 부분' {
			// 스왑
			temp := s.node[parent]
			s.node[parent] = s.node[index]
			s.node[index] = temp
		}
		index = parent
	}
}
// Poll 요소를 반환하고 삭제한다.
func (s *heap) Poll() int {
	result := s.node[1]
	//맨마지막 요소 루트노드로
	s.node[1] = s.node[len(s.node)-1]
	// 크기 줄이기
	s.node = append(s.node[0 : len(s.node)-1])
	parent := 1

	for parent*2 < len(s.node) {
		//왼쪽부터
		min := s.node[parent*2]
		minIndex := parent * 2
		//오른쪽이 더크면 변경
		
		if parent*2+1 < len(s.node) && '함수를 만들어서 관리할 부분' {
			min = s.node[parent*2+1]
			minIndex = parent*2 + 1
		}
		//부모가 더크면 종료
		if '함수를 만들어서 관리할 부분' {
			break
		}
		//스왑
		temp := s.node[minIndex]
		s.node[minIndex] = s.node[parent]
		s.node[parent] = temp
		//밑으로
		parent = minIndex
	}
	return result
}

Poll()과 Add()메서드를 잘보면 결국 분기조건에 따라 최대힙, 최소힙이 정해지는데 이분을 수정하면 효율적으로 heap을 구현할수 있다.

type heap struct {
	node  []int
	check func(a, b int) bool
}

func NewMaxHeap() *heap {
	return &heap{[]int{0}, func(x int, y int) bool { return x < y }}
}
func NewMinHeap() *heap {
	return &heap{[]int{0}, func(x int, y int) bool { return x > y }}
}

기본값은 maxheap으로 하고 익명 함수를 파라미터로 받고 싶었는데 Go에서는 오버로딩을 지원하지 않기 때문에 함수를 각각 나눠 분기조건으로 들어갈 함수를 익명 함수로 따로 구현한다.

func (s *heap) swap(i, j int) {
	s.node[i], s.node[j] = s.node[j], s.node[i]
}

그리고 swap 중복코드를 함수로 정리하고 제네릭을 적용한다.

적용해보기

11279.최대 힙

1927.최소 힙

전체 코드

heap_test.go

package heap

import (
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"

	"testing"
)

type TestSuite struct {
	suite.Suite
	h *heap[int]
	p *heap[int]
}

func TestRunSuit(t *testing.T) {
	suite.Run(t, new(TestSuite))
}

// 함수가 실행할때마다
func (ts *TestSuite) SetupTest() {
	ts.h = NewMaxHeap[int]()
	ts.p = NewMinHeap[int]()
}

func (ts *TestSuite) TestHeap_Add() {
	ts.h.Add(1)
	ts.h.Add(2)
	ts.h.Add(3)

	assert.Equal(ts.T(), 3, len(ts.h.node)-1, "maxHeap add error")

	ts.p.Add(1)
	ts.p.Add(2)
	ts.p.Add(3)

	assert.Equal(ts.T(), 3, len(ts.p.node)-1, "minHeap add error")

}
func (ts *TestSuite) TestHeap_Poll() {
	ts.h.Add(1)
	ts.h.Add(3)
	ts.h.Add(10)
	result := ts.h.Poll()

	assert.Equal(ts.T(), 10, result)

	ts.p.Add(1)
	ts.p.Add(3)
	ts.p.Add(10)
	result = ts.p.Poll()

	assert.Equal(ts.T(), 1, result)
}

func (ts *TestSuite) TestHeap_Peek() {
	ts.h.Add(1)
	ts.h.Add(100)
	result := ts.h.Peek()

	assert.Equal(ts.T(), 100, result, "Peek error")
	assert.Equal(ts.T(), 2, len(ts.h.node)-1, "size error")
}

func (ts *TestSuite) TestHeap_Size() {
	ts.h.Add(1)
	ts.h.Add(1)
	ts.h.Add(1)
	ts.h.Add(1)
	assert.Equal(ts.T(), 4, ts.h.Size(), "Size error")
}

func (ts *TestSuite) TestHeap_Empty() {
	ts.h.Add(1)
	assert.Equal(ts.T(), false, ts.h.Empty())
}

heap.go

package heap

type number interface {
	int8 | int16 | int32 | int64 | int
}
type heap[T any] struct {
	node  []T
	check func(a, b T) bool
}

func NewHeap[T any](f func(a, b T) bool) *heap[T] {
	return &heap[T]{make([]T, 1), f}
}
func NewMaxHeap[T number]() *heap[T] {
	return &heap[T]{make([]T, 1), func(x T, y T) bool { return x < y }}
}
func NewMinHeap[T number]() *heap[T] {
	return &heap[T]{make([]T, 1), func(x T, y T) bool { return x > y }}
}

// Add heap에 요소를 추가한다.
// num: 추가하고 싶은 값
func (s *heap[T]) Add(num T) {
	// 맨 마지막 인덱스에 요소 추가
	s.node = append(s.node, num)
	// 연산
	index := len(s.node) - 1
	for index > 1 {
		parent := index / 2

		if s.check(s.node[parent], s.node[index]) {
			// 스왑
			s.swap(parent, index)
			//temp := s.node[parent]
			//s.node[parent] = s.node[index]
			//s.node[index] = temp
		}
		index = parent
	}
}

// Poll 요소를 반환하고 삭제한다.
func (s *heap[T]) Poll() T {
	result := s.node[1]
	//맨마지막 요소 루트노드로
	s.node[1] = s.node[len(s.node)-1]
	// 크기 줄이기
	s.node = append(s.node[0 : len(s.node)-1])
	parent := 1

	for parent*2 < len(s.node) {
		//왼쪽부터
		min := s.node[parent*2]
		minIndex := parent * 2
		//오른쪽이 더크면 변경

		if parent*2+1 < len(s.node) && s.check(min, s.node[parent*2+1]) {
			min = s.node[parent*2+1]
			minIndex = parent*2 + 1
		}
		//부모가 더크면 종료
		if s.check(min, s.node[parent]) {
			break
		}
		//스왑
		s.swap(minIndex, parent)
		//temp := s.node[minIndex]
		//s.node[minIndex] = s.node[parent]
		//s.node[parent] = temp
		//밑으로
		parent = minIndex
	}
	return result
}

// Peek 최상 위 요소 반환
func (s *heap[T]) Peek() T {

	return s.node[1]

}

// Size heap 크기 반환
func (s *heap[T]) Size() int {
	return len(s.node) - 1
}

// Empty 공백 확인
func (s *heap[T]) Empty() bool {
	return len(s.node) == 1
}

func (s *heap[T]) swap(i, j int) {
	s.node[i], s.node[j] = s.node[j], s.node[i]
}