sync.Mutex探索 part1 - 算法探索 Treap

火莲华
0 阅读6分钟

前言

sync.Mutex 即是我们熟悉的互斥锁, 这个结构同时也用到了runtime.sema. 关于sema读者可能会觉得陌生, 因为这个数据结构并不是开放给开发者使用的, 不过它的设计值得专门花篇幅去介绍. 故关于这个话题, 这里将分为三部分进行记录, 第一部分介绍runtime.sema 用到的算法 Treap; 第二部分正式介绍runtime.sema的实现, 并结合第一部分谈谈go是怎么实现 Treap; 而第三部分再介绍主题 Mutex的底层实现.

Treap即 Tree + Heap, 这个数据结构的初衷是为了更有效率地构建一颗比较平衡的二叉搜索树而诞生. 就像前面介绍的, golang的底层源码也有对该数据结构的应用--runtime.sema.

这篇笔记将重点介绍这个算法, 并实现对应的插入, 删除, 查询操作.

参考文献:

courses.grainger.illinois.edu/cs473/sp201…

算法介绍

Treap 是在二叉搜索树的基础上杂糅了堆结构的特征, 使之形成类似完全二叉搜索树的效果.

特性:

  1. Treap的每个节点的key必须满足二叉搜索树的特性
  2. Treap的每个节点的priority必须满足堆的特性

如下图,为通过最小堆构建的Treap

image.png

其中红色节点为key, 绿色节点为priority. 这两个值决定了一个Treap的构成结构. 看到这里读者可能会疑惑:

  1. key是用来搜索的, priority是用来干嘛的?
  2. Treap是怎样来保证构建出来的树拟似完全二叉搜索树?

这两个问题其实涉及到Treap的本质了, 首先priority是为了保证Treap能让每个节点符合堆的特征, 除此之外没有其他意义. 而priority 是随机生成的, 也就是Treap其实就是通过随机数+堆的特性, 使整个树结构趋向于完全二叉搜索树.

我们将上图的 key 和 priority 分开来看:

对应的key组成一个二叉搜索树 image.png

对应的priority组成一个最小堆 image.png

场景实现

接下来我们来尝试实现这样一种数据结构.

数据结构

type Node[kType constraints.Ordered, vType any] struct {
	key      kType
	val      vType
	priority int
	left     *Node[kType, vType]
	right    *Node[kType, vType]
	parent   *Node[kType, vType] // 这里parent指向父节点, 方便实现递推调整
} // 这里节点的key应该是可排序的

type Tree[kType constraints.Ordered, vType any] struct {
	root *Node[kType, vType]
	r    *rand.Rand // 用于生成随机优先级
	size int
}

func (t *Tree[kType, vType]) Put(key kType, val vType)
func (t *Tree[kType, vType]) Get(key kType) (vType, bool)
func (t *Tree[kType, vType]) Del(key kType)

Node是Treap对应的节点, 而Tree是对应的实现.

旋转

在插入和删除过程中会涉及到旋转操作以保证二叉搜索树的特性. 所以我们先实现其左旋,右旋操作:

image.png 右旋, 本质上是将节点向右下移的同时不破坏树的完整性的操作
代码实现:

func (t *Tree[kType, vType]) rightRotate(n *Node[kType, vType]) {
    p := n.parent
    s := n.left
    sr := s.right
    // 左节点上移(对应图中的B)
    if p != nil {
            if p.left == n {
                    p.left = s
            } else {
                    p.right = s
            }
    }
    s.parent = p
    
    // 自节点向右下移(对应图中的A)
    s.right = n
    n.parent = s
    
    // 左节点的右节点调整到自节点的左节点(对应图中的E)
    n.left = sr
    if sr != nil {
            sr.parent = n
    }
}

image.png 左旋, 本质上是将节点向左下移的同时不破坏树的完整性的操作
代码实现:

func (t *Tree[kType, vType]) leftRotate(n *Node[kType, vType]) {
	p := n.parent
	s := n.right
	var sl *Node[kType, vType]
	sl = s.left
        
        // 右节点上移(对应图中的B)
	if p != nil {
		if p.left == n {
			p.left = s
		} else {
			p.right = s
		}
	}
	s.parent = p
        
        // 自节点向左下移(对应图中的A)
	s.left = n
        n.parent = s
        
        // 右节点的左节点调整到自节点的右节点(对应图中的C)
	n.right = sl
	if sl != nil {
		sl.parent = n
	}
}

细心的朋友可以观察到, 左旋和右旋其实是一对镜像操作

Get方法

下面我们先从最简单的Get方法开始实现:
get方法跟普通二叉搜索树的get逻辑是一致的

// find elem or privious elem
func (t *Tree[kType, vType]) get(key kType) (*Node[kType, vType], bool) {
    cur := t.root
    pre := cur
    for cur != nil {
        pre = cur
        if cur.key < key {
            cur = cur.right // 搜索key大于当前节点key, 向右搜索
            continue
        }

        if cur.key > key { // 搜索key小于当前节点key, 向左搜索
            cur = cur.left
            continue
        }

        return cur, true
    }

    return pre, false
}

func (t *Tree[kType, vType]) Get(key kType) (vType, bool) {
    n, ok := t.get(key)
    if !ok {
        return n.val, ok
    }

    var zero vType
    return zero, ok
}

这里封装了get方法, 如果找不到就返回对应位置的节点, 方便后面del, put复用

Put方法

插入逻辑.首先需要先找到对应二叉搜索树的位置, 再通过priority值向上进行堆调整.
这里我们在之前举的Treap例子里, 插入节点0, priority是40:

image.png

  1. 先找到0的位置在(1,134)节点的左节点. 然后发现134比40大, 不符合最小堆. 进行右旋调整.
  2. 经过右旋, (0,40)在(3,53)的左节点. 依然不符合最小堆. 再次向上调整, 进行右旋.
  3. 这时 (0,40)在(6,39)的左节点, 符合最小堆结构, 且整体符合二叉搜索树结构. 完成调整.

代码实现:

func (t *Tree[kType, vType]) Put(key kType, val vType) {
	prio := t.r.Int31()
	newNode := &Node[kType, vType]{
		key:      key,
		val:      val,
		priority: int(prio),
	}

	if t.root == nil {
		t.root = newNode
                return // 根节点直接返回
	}

	curNode, ok := t.get(key)
	if ok {
		curNode.val = val // 找到节点, 直接修改值
		return
	}

	newNode.parent = curNode
	if curNode.key > key {
		curNode.left = newNode
	} else {
		curNode.right = newNode
	} // 新节点放到对应位置

	cur := newNode
        // 开始向上调整
	for cur.parent != nil && cur.priority < cur.parent.priority {
		isRoot := cur.parent == t.root
		switch cur {
		case cur.parent.left:
			t.rightRotate(cur.parent) // 处于左节点, 通过右旋向上调整
		case cur.parent.right:
			t.leftRotate(cur.parent) // 处于右节点, 通过左旋向上调整
		}

		if isRoot {
			t.root = cur // 防止丢失根节点
		}
	}

	t.size++
}

Del方法

删除逻辑, 先找到待删除的节点, 让该节点向下调整, 直到为叶子节点, 直接删除. 此时可假设该节点priority无限大.

如下图, 同样的案例, 这次将之前添加的(0,40)节点删除. image.png

  1. 首先先将待删除节点的priority置为无限大(图中用NaN来标识), 并右旋往下调整
  2. 此时(0, NaN) 在(3,53)节点的右节点, 继续右旋向下调整
  3. 此时(0, NaN) 在(5, 201)节点右边, 且为叶子节点.
  4. 删除节点, 删除操作完成

代码实现:

func (t *Tree[kType, vType]) Del(key kType) {
	n, ok := t.get(key)
	if !ok {
		return // 没找到, 中断
	}

	for n.left != nil || n.right != nil { // 不是叶子节点, 向下调整
		leftPrio, rightPrio := -1, -1
		if n.left != nil {
			leftPrio = n.left.priority
		}

		if n.right != nil {
			rightPrio = n.right.priority
		}

		isRoot := t.root == n
		switch {
		case leftPrio == -1: // 只有右节点, 左旋调整
			t.leftRotate(n)
		case rightPrio == -1: // 只有左节点, 右旋调整
			t.rightRotate(n)
		case leftPrio > rightPrio: // 左右节点都存在, 将较小权重的节点向上移
			t.leftRotate(n)
		default:
			t.rightRotate(n)
		}

		if isRoot {
			t.root = n.parent // 避免丢失根节点
		}

	}

	if t.root == n {
		t.root = nil // 只有根节点且呆删节点为根节点, 直接删除
		return
	}

	if n.parent.left == n {
		n.parent.left = nil
	} else {
		n.parent.right = nil
	}
}

后记

通过前面的介绍得知, Treap 本质上还是高度依赖于随机数. 故当各节点生成的priority 越随机, Treap就越接近完全二叉搜索树. 不过由于更新时并不需要调整整棵树, 故显而易见的优点是只要牺牲一点查询性能, 便使其更新的效率远高于完全二叉搜索树. 而且易于实现, 比较适用于需要频繁更新的场景.

记录于 2026.03.31

avatar
文章
阅读
粉丝