[Data Structure] Tree(5) - Red-Black Tree Implementation

이번 포스팅에선 "Red-Black Tree"를 Swift로 구현할 것이다.

만약 개념에 대한 내용이 궁금하다면 해당 포스팅을 참고 바란다. 

 

자세한 구현은 아래 링크에서 확인해 볼 수 있다. 

https://github.com/jungseokyoung-cloud/Data-Structure/tree/main/Data-Structure/Tree.playground/Sources/Red-BlackTree

Data-Structure/Data-Structure/Tree.playground/Sources/Red-BlackTree at main · jungseokyoung-cloud/Data-Structure

 

 

RedBlack Node

우선, RedBlackNode는 다음과 같이 구현해 준다. 

final class RedBlackNode<T: Comparable>: Node<T> {
	var isRed: Bool = true
	
	required init(value: T) {
		super.init(value: value)
	}
	
	required init(value: T, parent: Node<T>?) {
		super.init(value: value, parent: parent)
	}
}

색깔을 구분하기 위한 isRed 프로퍼티를 추가한다.

 

RedBlack Tree 

RB 트리에서 NIL노드는 별도의 메모리를 할당하지 않고, 단순히 nil로 지정했다.

 

insert(_:)

public override func insert(_ value: T) {
	guard let node = super.performInsert(value) as? RedBlackNode<T> else { return }
	
	insertBalancing(for: node)
}

insert(_:) 메서드는 insert수행한 후, insertBalancing(for:) 메서드를 수행한다.

 

 

insertBalancing(for:)

@discardableResult
func insertBalancing(for node: RedBlackNode<T>) -> RedBlackNode<T>? {
	guard node !== root else {
		node.isRed = false
		return node
	}
	
	// 연속해서 빨간색이 나오지 않는 경우, 종료
	guard
		let parent = node.parent as? RedBlackNode<T>,
		parent.isRed
	else { return node }
	
	if let uncle = node.uncle as? RedBlackNode<T>, uncle.isRed {
		return recoloring(for: node)
	} else {
		return restructuring(for: node)
	}
}

insertBalancing(for:)메서드는 restructuring을 할지 recoloring을 할지 결정하는 메서드이다. 

 

우선, 부모노드가 빨간색인 경우 #4 조건을 위반하게 되어, self-balancing이 필요하다. 

#4: 빨간색 노드의 자식은 검은색이다.  (즉, 2 연속으로 빨간색이 나올 수 없다.)

 

// 연속해서 빨간색이 나오지 않는 경우, 종료
guard
	let parent = node.parent as? RedBlackNode<T>,
	parent.isRed
else { return node }

그렇기 때문에 다음과 같이 부모 노드가 검은색일 경우, 다음과 같이 종료한다.

 

if let uncle = node.uncle as? RedBlackNode<T>, uncle.isRed {
	return recoloring(for: node)
} else {
	return restructuring(for: node)
}

다음으로, 삼촌 노드가 빨간색일 경우 recoloring을 수행하고, 

검은색일 경우, restructing을 수행한다. 

 

restructuring(for:)

@discardableResult
func restructuring(for node: RedBlackNode<T>) -> RedBlackNode<T>? {
	guard
		let parent = node.parent as? RedBlackNode<T>,
		let grandParent = parent.parent as? RedBlackNode<T>
	else { return node }
	grandParent.isRed = true

	// LL Case
	if node.isLeftChild && parent.isLeftChild {
		parent.isRed = false
		return rightRotate(for: grandParent, pivot: parent)
		
	// RR Case
	} else if node.isRightChild && parent.isRightChild {
		parent.isRed = false
		return leftRotate(for: grandParent, pivot: parent)
		
	// LR Case
	} else if node.isRightChild && parent.isLeftChild {
		node.isRed = false
		let pivot = leftRotate(for: parent, pivot: node)
		return rightRotate(for: grandParent, pivot: pivot)
		
	// RL Case
	} else {
		node.isRed = false
		let pivot = rightRotate(for: parent, pivot: node)
		return leftRotate(for: grandParent, pivot: pivot)
	}
}

RR Case, LL Case의 경우 부모 노드와 조부모 노드의 색상을 변경해 주고, Rotation 작업을 수행한다. 

또한, RL Case, LR Case의 경우 새로운 노드와 조부모 노드의 색상을 변경한 후, Rotation 작업을 수행한다.

 

 

이때, NIL노드가 삼촌노드일수도 있기 때문에, 다음과 같이 optional 바인딩을 진행해 준다.

 

recoloring(for:)

@discardableResult
func recoloring(for node: RedBlackNode<T>) -> RedBlackNode<T>? {
	guard
		let parent = node.parent as? RedBlackNode<T>,
		let grandParent = parent.parent as? RedBlackNode<T>,
		let uncle = node.uncle as? RedBlackNode<T>
	else { return node }
	
	grandParent.isRed = true
	parent.isRed = false
	uncle.isRed = false
	
	return insertBalancing(for: grandParent)
}

recoloring(for:)메서드는 다음과 같이 부모노드와, 삼촌 노드를 검은색으로 변경한다. 

이후, 조부모 노드를 빨간색으로 변경한다. 

 

하지만, 조부모 노드에서 다시 #4조건을 위반할 수 있기 때문에, 재귀적으로 insertBalancing(for:)메서드를 수행해 준다.

 

 

remove(_:)

public override func remove(_ value: T) -> T? {
	guard let removeNode = search(value) else { return nil }

	let removedReturnType = super.performRemove(node: removeNode)
	removeBalancing(
		removedNode: removedReturnType.removedNode,
		replaceNode: removedReturnType.replaceNode,
		parentNode: removedReturnType.parentNode
	)
	
	return removeNode.value
}

remove(_:)메서드는 removeBalancing(removedNode:replaceNode:parentNode:)를 통해 균형을 잡는다. 

 

removeBalancing(removedNode:replaceNode:parentNode:)

func removeBalancing(
	removedNode: Node<T>?,
	replaceNode: Node<T>?,
	parentNode: Node<T>?
) {
	guard let node = removedNode as? RedBlackNode<T>, !node.isRed else { return }
	
	// 삭제된 노드가 검정색이라면, Extra-Black을 부여
	removeExtraBlack(parent: parentNode, replaceNode: replaceNode)
	
	(root as? RedBlackNode<T>)?.isRed = false
}

removeBalancing은 삭제되는 노드가 검은색이라면, 문제가 발생한다. 따라서, guard문을 통해 삭제되는 노드가 검은색인지 확인한다. 

 

다음으로, removeExtraBlack메서드를 통해 ExtraBlack을 제거한다. 

 

removeExtraBlack(parent:replaceNode:)

func removeExtraBlack(parent: Node<T>?, replaceNode: Node<T>?) {
	// Red-And-Black Node
	if let replaceNode = replaceNode as? RedBlackNode<T>, replaceNode.isRed {
		replaceNode.isRed = false
	} else {
		removeDoublyBlack(parent: parent, doublyBlack: replaceNode)
	}
}

이때, replaceNode는 삭제된 노드의 위치에 위치한 노드이다.

replaceNode가 빨간색이라면, Red-And-Black이므로 검은색으로 변경해 주면 되지만, 

Double-Black이라면 여러 분기로 나뉘기 때문에, removeDoublyBlack메서드를 통해 제거한다.

 

removeDoublyBlack(parent:doublyBlack:)

func removeDoublyBlack(parent: Node<T>?, doublyBlack: Node<T>?) {
	guard
		let parent = parent as? RedBlackNode<T>,
		let sibling = (parent.left === doublyBlack ? parent.right : parent.left) as? RedBlackNode<T>
	else { return }
	
	if !sibling.isRed {
		if
			let sibilingChild = (sibling.isLeftChild ? sibling.left : sibling.right) as? RedBlackNode<T>,
			sibilingChild.isRed {
			removeCase4(
				parent: parent,
				sibling: sibling,
				child: sibilingChild
			)
		} else if
			let sibilingChild = (sibling.isLeftChild ? sibling.right : sibling.left) as? RedBlackNode<T>,
			sibilingChild.isRed {
			removeCase3(
				parent: parent,
				sibling: sibling,
				child: sibilingChild
			)
		} else {
			removeCase2(parent: parent, sbiling: sibling)
		}
	} else {
		removeCase1(parent: parent, sibling: sibling, node: doublyBlack as? RedBlackNode<T>)
	}
}

분기가 복잡해 보이지만 살펴보면 단순하다. 

 

우선, 형제노드의 색깔이 검은색이고, 형제노드와 같은 위치의 형제노드의 자식(LL 혹은 RR)이 빨간색인 경우,

if !sibling.isRed {
	if
		let sibilingChild = (sibling.isLeftChild ? sibling.left : sibling.right) as? RedBlackNode<T>,
		sibilingChild.isRed {
		removeCase4(
			parent: parent,
			sibling: sibling,
			child: sibilingChild
		)
	}

case4에 해당되기 때문에, removeCase4메서드를 통해 Doubly-Black을 제거해 준다. 

 

다음으로, 형제노드의 색깔이 검은색이고, 형제노드와는 반대편 위치에 존재하는 형제노드의 자식(LR 혹은 RL)이 빨간색이고, 

같은 위치의 형제노드는(LL 혹은 RR)이 검은색인 경우,

if !sibling.isRed {
	...
	} else if
		let sibilingChild = (sibling.isLeftChild ? sibling.right : sibling.left) as? RedBlackNode<T>,
		sibilingChild.isRed {
		removeCase3(
			parent: parent,
			sibling: sibling,
			child: sibilingChild
		)
	}

 

case3에 해당하여 removeCase3메서드를 통해 Doubly-Black을 제거해 준다.

 

만약, 형제노드가 검은색이고, 형제노드의 두 자식 다 검은색이라면, 

if !sibling.isRed {
	...
	} else {
		removeCase2(parent: parent, sbiling: sibling)
	}
}

case2이므로, removeCase2메서드를 통해 Doubly-Black을 제거해 준다.

 

마지막으로 형제노드가 빨간색인 경우는 case1에 해당되므로 removeCase1메서드를 호출한다.

else {
	removeCase1(parent: parent, sibling: sibling, node: doublyBlack as? RedBlackNode<T>)
}

 

removeCase4(parent:sibling:child:)

func removeCase4(
	parent: RedBlackNode<T>,
	sibling: RedBlackNode<T>,
	child: RedBlackNode<T>
) {
	sibling.isRed = parent.isRed
	
	child.isRed = false
	parent.isRed = false
	
	if sibling.isLeftChild {
		rightRotate(for: parent, pivot: sibling)
	} else {
		leftRotate(for: parent, pivot: sibling)
	}
}

형제노드는 부모노드의 색상으로 변경하고, 형제노드의 자식노드와 부모노드를 검은색으로 변경한다. 

이후, 각 상황에 맞게 rotation을 수행한다. 

 

removeCase3(parent:sbling:child:)

func removeCase3(
	parent: RedBlackNode<T>,
	sibling: RedBlackNode<T>,
	child: RedBlackNode<T>
) {
	sibling.isRed = true
	child.isRed = false
	
	if sibling.isLeftChild {
		leftRotate(for: sibling, pivot: child)
	} else {
		rightRotate(for: sibling, pivot: child)
	}
	
	guard
		let newChild = (sibling.isLeftChild ? child.left : child.right) as? RedBlackNode<T>
	else { return }
	
	removeCase4(parent: parent, sibling: child, child: newChild)
}

해당 경우에는 형제노드를 빨간색으로, 형제의 자식노드를 검은색으로 변경한다. 

이후, Rotation 작업을 진행해 주게 되면 case4가 되기 때문에 removeCase4를 호출해 준다.

 

removeCase2(parent:sibling:)

func removeCase2(parent: RedBlackNode<T>, sibling: RedBlackNode<T>) {
	sibling.isRed.toggle()
	// 부모에게 Extra-Black을 위임
	removeExtraBlack(parent: parent.parent, replaceNode: parent)
}

형제노드의 색상을 변경해 주고, Extra-Black을 부모에게 넘겨준다. 

이후 부모노드에게 ExtraBlack을 위임한다.

 

removeCase1(parent:sibling:node:)

func removeCase1(parent: RedBlackNode<T>, sibling: RedBlackNode<T>, node: RedBlackNode<T>?) {
	parent.isRed = true
	sibling.isRed = false
	
	if sibling.isLeftChild {
		rightRotate(for: parent, pivot: sibling)
	} else {
		leftRotate(for: parent, pivot: sibling)
	}
	
	removeDoublyBlack(parent: parent, doublyBlack: node)
}

case1의 경우에는 부모 노드와 형제노드의 색상을 변경한 후, rotation 작업을 수행한다. 

이후 case2, case3, case4에 맞춰서 DoublyBlack을 해결하면 된다.