• red-black tree와 유사하지만 경우의 수가 더 적은 2-3 트리와 일대일 대응이 되는 self-balancing binary tree를 구현해보았다.

• red-black tree는 아주 아주 오래 걸릴 것 같아서 더 쉬운 버전을 구현한다고 했는데 하루 종일 걸렸다.

  • red-black tree보다는 간단해서 순한 맛인데,

  • 트리 회전에 대해서 생각하지 않고 순수 함수를 이용한 순한 맛인데

  • 타입스크립트의 표현력 깊은 타입 시스템까지 활용해서 분명히 순한 맛인데

  • 어려웠다.

    • 고려할 경우의 수가 많았다.

  • 테스트도 제대로 못 해보았지만 구현한 것에 의미를 두고 잘 동작하길 기원만 하고 넘어가기로 했다. 부끄럽다.

     

// Copyright 2024-2024 Dongook Lee
// SPDX-License-Identifier: MIT

// AA-like tree implementation with one tag bit per node
// https://en.wikipedia.org/wiki/AA_tree

type Tree = null | Node2 | Node3

// A tag value of true corresponds to black nodes in a red-black tree
type Node2 = [tag: true, l: Tree, x: number, r: Tree]
type Node3 = [tag: true, l: Tree, x: number, [tag: false, m: Tree, y: number, r: Tree]]

function isNode2(tree: Tree): tree is Node2 {
    return tree !== null && (tree[3] === null || tree[3][0])
}

const node2 = (l: Tree, x: number, r: Tree): Node2 => [true, l, x, r]
const replaceLeft2 = ([, _l, x, r]: Node2, l2: Tree): Node2 => node2(l2, x, r)
const replaceValue2 = ([, l, _x, r]: Node2, x2: number): Node2 => node2(l, x2, r)
const replaceRight2 = ([, l, x, _r]: Node2, r2: Tree): Node2 => node2(l, x, r2)

const node3 = (l: Tree, x: number, m: Tree, y: number, r: Tree): Node3 => [true, l, x, [false, m, y, r]]
const replaceLeft3 = ([, _l, x, [, m, y, r]]: Node3, l2: Tree): Node3 => node3(l2, x, m, y, r)
const replaceLeftValue3 = ([, l, _x, [, m, y, r]]: Node3, x2: number): Node3 => node3(l, x2, m, y, r)
const replaceMiddle3 = ([, l, x, [, _m, y, r]]: Node3, m2: Tree): Node3 => node3(l, x, m2, y, r)
const replaceRightValue3 = ([, l, x, [, m, _y, r]]: Node3, y2: number): Node3 => node3(l, x, m, y2, r)
const replaceRight3 = ([, l, x, [, m, y, _r]]: Node3, r2: Tree): Node3 => node3(l, x, m, y, r2)

// Insert
const upgradeLeft2 = ([, _l, x, r]: Node2, [, ll, y, lr]: Node2): Node3 => node3(ll, y, lr, x, r)
const upgradeRight2 = ([, l, x, _r]: Node2, [, rl, y, rr]: Node2): Node3 => node3(l, x, rl, y, rr)

function upgradeLeft3([, _l, x, [, m, y, r]]: Node3, [, ll, z, lr]: Node2): Node2 {
    return node2(node2(ll, z, lr), x, node2(m, y, r))
}
function upgradeMiddle3([, l, x, [, _m, y, r]]: Node3, [, ml, z, mr]: Node2): Node2 {
    return node2(node2(l, x, ml), z, node2(mr, y, r))
}
function upgradeRight3([, l, x, [, m, y, _r]]: Node3, [, rl, z, rr]: Node2): Node2 {
    return node2(node2(l, x, m), y, node2(rl, z, rr))
}

type InsertionResult =
    | [upgraded: false, newNode: Tree]
    | [upgraded: true, newNode: Node2]

function insert(tree: Tree, v: number): InsertionResult {
    if (tree === null) {
        return [true, node2(null, v, null)]
    }
    if (isNode2(tree)) {
        const [, l, x, r] = tree
        if (v >= x) {
            const [up, r2] = insert(r, v)
            return up ? [false, upgradeRight2(tree, r2)] : [false, replaceRight2(tree, r2)]
        }
        const [up, l2] = insert(l, v)
        return up ? [false, upgradeLeft2(tree, l2)] : [false, replaceLeft2(tree, l2)]
    }
    // Node3
    const [, l, x, [, m, y, r]] = tree
    if (v >= y) {
        const [up, r2] = insert(r, v)
        return up ? [true, upgradeRight3(tree, r2)] : [false, replaceRight3(tree, r2)]
    }
    if (v >= x) {
        const [up, m2] = insert(m, v)
        return up ? [true, upgradeMiddle3(tree, m2)] : [false, replaceMiddle3(tree, m2)]
    }
    const [up, l2] = insert(l, v)
    return up ? [true, upgradeLeft3(tree, l2)] : [false, replaceLeft3(tree, l2)]
}

type DeletionResult =
    | [downgraded: false, newNode: Tree]
    | [downgraded: true, newNode: Node3 | null]

function downgradeLeft2([, _l, x, r]: Node2, l2: Node3 | null): DeletionResult {
    if (r === null) {
        throw new TypeError("expected non-null sibling, got null")
    }
    if (isNode2(r)) {
        const [, rl, y, rr] = r
        return [true, node3(l2, x, rl, y, rr)]
    }
    const [, rl, y, [, rm, z, rr]] = r
    return [false, node2(node2(l2, x, rl), y, node2(rm, z, rr))]
}

function downgradeRight2([, l, x, _r]: Node2, r2: Node3 | null): DeletionResult {
    if (l === null) {
        throw new TypeError("expected non-null sibling, got null")
    }
    if (isNode2(l)) {
        const [, ll, y, lr] = l
        return [true, node3(ll, y, lr, x, r2)]
    }
    const [, ll, y, [, lm, z, lr]] = l
    return [false, node2(node2(ll, y, lm), z, node2(lr, x, r2))]
}

function downgradeLeft3([, _l, x, [, m, y, r]]: Node3, l2: Node3 | null): DeletionResult {
    if (m === null) {
        throw new TypeError("expected non-null sibling, got null")
    }
    if (isNode2(m)) {
        const [, ml, z, mr] = m
        return [false, node2(node3(l2, x, ml, z, mr), y, r)]
    }
    const [, ml, z, [, mm, w, mr]] = m
    return [false, node3(node2(l2, x, ml), z, node2(mm, w, mr), y, r)]
}

function downgradeMiddle3([, l, x, [, _m, y, r]]: Node3, m2: Node3 | null): DeletionResult {
    if (r === null) {
        throw new TypeError("expected non-null sibling, got null")
    }
    if (isNode2(r)) {
        const [, rl, z, rr] = r
        return [false, node2(l, x, node3(m2, y, rl, z, rr))]
    }
    const [, rl, z, [, rm, w, rr]] = r
    return [false, node2(l, x, node2(node2(m2, y, rl), z, node2(rm, w, rr)))]
}

function downgradeRight3([, l, x, [, m, y, _r]]: Node3, r2: Node3 | null): DeletionResult {
    if (m === null) {
        throw new TypeError("expected non-null sibling, got null")
    }
    if (isNode2(m)) {
        const [, ml, z, mr] = m
        return [false, node2(l, x, node3(ml, z, mr, y, r2))]
    }
    const [, ml, z, [, mm, w, mr]] = m
    return [false, node2(l, x, node2(node2(ml, z, mm), w, node2(mr, y, r2)))]
}

function isLeafNode(node: Node2 | Node3): boolean {
    return node[1] === null
}

function popMinValue(tree: Tree): [number, DeletionResult] {
    if (tree === null) {
        throw new Error("expected non-empty tree, got null")
    }
    if (isLeafNode(tree)) {
        if (isNode2(tree)) {
            const [, _l, x, _r] = tree
            return [x, [true, null]]
        }
        const [, _l, x, [, _m, y, _r]] = tree
        return [x, [false, node2(null, y, null)]]
    }
    if (isNode2(tree)) {
        const [, l] = tree
        const [x2, [down, l2]] = popMinValue(l)
        return down ? [x2, downgradeLeft2(tree, l2)] : [x2, [false, replaceLeft2(tree, l2)]]
    }
    const [, l] = tree
    const [x2, [down, l2]] = popMinValue(l)
    return down ? [x2, downgradeLeft3(tree, l2)] : [x2, [false, replaceLeft3(tree, l2)]]
}

function remove(tree: Tree, v: number): DeletionResult {
    if (tree === null) {
        throw new Error("value not found")
    }
    if (isLeafNode(tree)) {
        if (isNode2(tree)) {
            const [, _l, x, _r] = tree
            if (v === x) {
                return [true, null]
            }
            throw new Error("value not found")
        }
        const [, _l, x, [, _m, y, _r]] = tree
        if (v === x) {
            return [false, node2(null, y, null)]
        }
        if (v === y) {
            return [false, node2(null, x, null)]
        }
        throw new Error("value not found")
    }
    if (isNode2(tree)) {
        const [, l, x, r] = tree
        if (v > x) {
            const [down, r2] = remove(r, v)
            return down ? downgradeRight2(tree, r2) : [false, replaceRight2(tree, r2)]
        }
        if (v === x) {
            const [x2, [down, r2]] = popMinValue(r)
            return down ? downgradeRight2(replaceValue2(tree, x2), r2) : [false, node2(l, x2, r2)]
        }
        const [down, l2] = remove(l, v)
        return down ? downgradeLeft2(tree, l2) : [false, replaceLeft2(tree, l2)]
    }
    const [, l, x, [, m, y, r]] = tree
    if (v > y) {
        const [down, r2] = remove(r, v)
        return down ? downgradeRight3(tree, r2) : [false, replaceRight3(tree, r2)]
    }
    if (v === y) {
        const [y2, [down, r2]] = popMinValue(r)
        return down ? downgradeRight3(replaceRightValue3(tree, y2), r2) : [false, node3(l, x, m, y2, r2)]
    }
    if (v > x) {
        const [down, m2] = remove(m, v)
        return down ? downgradeMiddle3(tree, m2) : [false, replaceMiddle3(tree, m2)]
    }
    if (v === x) {
        const [x2, [down, m2]] = popMinValue(m)
        return down ? downgradeMiddle3(replaceLeftValue3(tree, x2), m2) : [false, node3(l, x2, m2, y, r)]
    }
    const [down, l2] = remove(l, v)
    return down ? downgradeLeft3(tree, l2) : [false, replaceLeft3(tree, l2)]
}

export function toInserted(t: Tree, v: number): Tree {
    return insert(t, v)[1]
}

export function toRemoved(t: Tree, v: number): Tree {
    return remove(t, v)[1]
}