用旋转 Treap 维护名次树

Posted on | In

Treap 是一种平衡树，通过在 BST 的基础上对每个结点钦定一个随机的优先级并在保证对于 BST 的性质不变的情况下，优先级满足堆的性质。Treap 分为两种，旋转 Treap（通过左旋/右旋操作维护性质）或者非旋 Treap（通过分裂/合并的方式维护性质），这里主要讲解旋转 Treap。

维护方式

如果用旋转 Treap 维护名次树，每个结点需要维护以下信息：

  • 两个儿子（如果不存在，记为空结点）
  • 权值
  • 出现次数
  • 子树大小
  • 优先级

同时还要维护整棵树的树根。

1
2
3
4
5
6
struct node {
  int l, r, cnt, val;
  long long pri;
  int sze;
} tree[MAXN];
int size, root;

旋转

如果我们按照 BST 的方式插入，因为新加入的结点一定是叶结点，难免会出现优先级不满足堆的性质的情况，所以我们要使用旋转的方式来在中序遍历不变的情况下对调某个结点和其父结点的父子关系。

具体解释如下图：

image-20210204234541058

左旋同理。易知旋转操作不改变整棵树的中序遍历（即不破坏 BST 性质）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
void rrotate(int &k) {
  int tmp = tree[k].l;
  tree[k].l = tree[tmp].r;
  tree[tmp].r = k;
  tree[k].sze = tree[tree[k].l].sze + tree[tree[k].r].sze + tree[k].cnt;
  tree[tmp].sze = tree[tree[tmp].l].sze + tree[tree[tmp].r].sze + tree[tmp].cnt;
  k = tmp; // 因为一般我们希
}

void lrotate(int &k) {
  int tmp = tree[k].r;
  tree[k].r = tree[tmp].l;
  tree[tmp].l = k;
  tree[k].sze = tree[tree[k].l].sze + tree[tree[k].r].sze + tree[k].cnt;
  tree[tmp].sze = tree[tree[tmp].l].sze + tree[tree[tmp].r].sze + tree[tmp].cnt;
  k = tmp;
}

插入

有了旋转，插入就很容易了。

  • 如果要插入到空树，直接新建结点。（无需维护父结点信息）
  • 如果插入的权值等于树根权值，出现次数 ++ 即可。
  • 如果以上条件皆不满足，根据权值与树根权值的关系判断应递归到左子树还是右子树继续插入（接下来以左子树为例继续讲解）。左子树插入完毕后（注意这里需要更新左儿子），如果左儿子的优先级与父亲优先级不满足堆性质，则将左儿子右旋至根。

经过以上条件，可以大致明白：新插入的结点会从叶结点一路向上旋转，直到自己的优先级与其目前父亲的优先级满足堆性质。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
void insert(int &k, int val) {
  if (k == 0) {
    tree[++size] = {0, 0, 1, val, static_cast<long long>(random()), 1};
    k = size;
    return;
  }
  if (tree[k].val == val) {
    tree[k].cnt++;
    tree[k].sze++;
  } else if (tree[k].val < val) {
    tree[k].sze++;
    insert(tree[k].r, val);
    if (tree[tree[k].r].pri < tree[k].pri) {
      lrotate(k);
    }
  } else {
    tree[k].sze++;
    insert(tree[k].l, val);
    if (tree[tree[k].l].pri < tree[k].pri) {
      rrotate(k);
    }
  }
}

删除

  • 如果树根的权值等于要删除的值：
    • 如果树根出现次数 $>1$，直接出现次数 — 即可。
    • 如果树根是叶子结点，将树置为空树即可。
    • 如果树根只有一个叶结点，将树置为其唯一的一棵子树即可。
    • 否则，比较左儿子与右儿子的优先级（以下假设左儿子优先级大于右儿子），并将左儿子右旋到根。接着递归处理当前树根的右儿子（即原先的树根，它的权值等于要删除的值）。
  • 否则，比较树根的权值与要删除的值，递归处理。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
void erase(int &k, int val) {
  if (k == 0) { return; }
  if (tree[k].val == val) {
    if (tree[k].cnt > 1) {
      tree[k].cnt--;
      tree[k].sze--;
      return;
    }
    if (tree[k].l == 0 && tree[k].r == 0) {
      k = 0;
      return;
    }
    if (tree[k].l != 0 && tree[k].r == 0) {
      k = tree[k].l;
      return;
    }
    if (tree[k].l == 0 && tree[k].r != 0) {
      k = tree[k].r;
      return;
    }
    if (tree[tree[k].l].pri < tree[tree[k].r].pri) {
      rrotate(k);
      erase(k, val);
      return;
    } else {
      lrotate(k);
      erase(k, val);
      return;
    }
  }
  if (tree[k].val > val) {
    tree[k].sze--;
    erase(tree[k].l, val);
    return;
  } else {
    tree[k].sze--;
    erase(tree[k].r, val);
    return;
  }
}

排名

利用子树大小判断该递归还是该返回即可。

1
2
3
4
5
6
7
8
9
10
11
int query_rnk(int k, int val) {
  if (k == 0) { return 0; } // 如果找不到
  if (tree[k].val == val) {
    return tree[tree[k].l].sze + 1;
  }
  if (tree[k].val > val) {
    return query_rnk(tree[k].l, val);
  } else {
    return tree[tree[k].l].sze + tree[k].cnt + query_rnk(tree[k].r, val);
  }
}

第 k 大

与排名类似。

1
2
3
4
5
6
7
8
9
10
int query_nth(int k, int val) {
  if (k == 0) { return 0; }
  if (val <= tree[tree[k].l].sze) {
    return query_nth(tree[k].l, val);
  } else if (val <= tree[tree[k].l].sze + tree[k].cnt) {
    return tree[k].val;
  } else {
    return query_nth(tree[k].r, val - tree[tree[k].l].sze - tree[k].cnt);
  }
}

前驱/后继

可以用前四个函数实现（插入，查询排名，查询第 k 大，删除），不过常数太大。

以前驱为例，一个结点的前驱一定是它作为右儿子的爸爸或者它左子树中序遍历的最后一个结点（即左儿子的右儿子的右儿子的右儿子……），所以：

如果还没经过目标结点：

Reference