黑匣子

GitHub跳转原题关系图 Canvas 关系图返回列表

两个堆维护已输出排名左侧与右侧元素，使每次 GET 的目标值位于右堆顶。

OJ: luogu

题目 ID: P1801

难度:普及+/提高

标签:双堆第k小heapqpython

日期: 2026-07-16 21:00

题意

元素按顺序加入；第 i 次 GET 输出当前集合第 i 小值。

思路

每次 GET 都要当前第 $i$ 小，直接对集合重新排序为什么不行？

集合最多 $2\times10^5$ 个数，GET 也最多 $2\times10^5$ 次，每次重排要 $O(m\log m)$ ，总代价 $O(m^2\log m)$ 必然超时。但题面有一个被忽略的细节： $i$ 每次只增加 $1$ 。相邻两次查询之间，目标只是从第 $i$ 小挪到第 $i+1$ 小，整个有序序列并没有大变动。

目标排名只挪一位，需要保存的信息能否缩小到"分界线"附近？

第 $i$ 小只取决于集合被切成"最小的 $i-1$ 个"和"其余"两半：答案就是右半的最小值。维护这条分界线只需要两侧各一个信息——左半的最大值（用来判断新数是否属于左半）、右半的最小值（答案本身）。取最大值和最小值正是堆的职责：左半用最大堆，右半用最小堆，这就是"双堆"的来历。

新元素进来，怎么决定它属于哪一半？

把新数和左半的最大值比较：若比它还小，说明新数应该算进"已输出的前 $i-1$ 小"，于是把它塞进左堆，再把左堆顶（最大的那个）挤进右堆；否则它本来就属于右半，直接进右堆。这样"左堆恰好是前 $i-1$ 小"的不变式始终保持。

GET 之后 $i$ 变成 $i+1$ ，两堆怎么推进？

这次输出的右堆顶就是第 $i$ 小，下一次它应当成为左半的一员。所以输出后把它移入左堆，右堆顶自然成为第 $i+1$ 小。每个元素进堆、出堆常数次，总复杂度降为 $O((m+n)\log m)$ 。

lower 用负数最大堆保存已经越过的 $i-1$ 个最小元素，upper 保存其余元素，小根堆顶就是下一答案。新值按分界线进入对应堆；每次输出 upper[0] 后，把它移入 lower，为下一排名作准备。

也可以换一个更直接的模型：用一个维护有序集合的平衡树（Treap），插入后直接用 kth 查询第 $i$ 小（见下面平衡树代码）。双堆版本的优点是只用两个堆，不需要完整有序结构。

Python 知识

负数把 heapq 变成最大堆。
heappushpop 一次完成插入和弹出，比两个独立操作更直接。
values[inserted:amount] 对应两次 GET 间新增的连续输入段。

代码

双堆解法（C++）

cpp

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
41
42
43
44
45
46#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int m, n;
    cin >> m >> n;
    vector<int> a(m);
    for (int i = 0; i < m; i++) {
        cin >> a[i];
    }
    vector<int> u(n);
    for (int i = 0; i < n; i++) {
        cin >> u[i];
    }

    priority_queue<int> L; // 最大堆，存储前 i-1 个最小元素
    priority_queue<int, vector<int>, greater<int>> R; // 最小堆，存储剩余元素，堆顶为第 i 小

    int inserted = 0;
    for (int q : u) {
        // 插入 a[inserted .. q-1]
        for (int i = inserted; i < q; i++) {
            int x = a[i];
            if (!L.empty() && x < L.top()) {
                L.push(x);
                R.push(L.top());
                L.pop();
            } else {
                R.push(x);
            }
        }
        inserted = q;
        // GET 操作
        cout << R.top() << "\n";
        L.push(R.top());
        R.pop();
    }

    return 0;
}

双堆解法（Python）

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24import heapq
import sys


data = iter(map(int, sys.stdin.buffer.read().split()))
value_count, query_count = next(data), next(data)
values = [next(data) for _ in range(value_count)]
queries = [next(data) for _ in range(query_count)]
lower = []
upper = []
inserted = 0
answers = []

for amount in queries:
    for value in values[inserted:amount]:
        if lower and value < -lower[0]:
            heapq.heappush(upper, -heapq.heappushpop(lower, -value))
        else:
            heapq.heappush(upper, value)
    inserted = amount
    answers.append(str(upper[0]))
    heapq.heappush(lower, -heapq.heappop(upper))

print("\n".join(answers))

平衡树解法（C++）

cpp

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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136#include <bits/stdc++.h>
using namespace std;

struct Treap {
    struct Node {
        int child[2] = {0, 0};
        int value = 0;
        int priority = 0;
        int count = 0;
        int size = 0;
    };

    vector<Node> tree;
    int root = 0;
    mt19937 rng;

    Treap(int max_nodes = 0) : rng(712367821) {
        tree.reserve(max_nodes + 1);
        tree.push_back(Node()); // node 0 is the null sentinel.
    }

    int new_node(int value) {
        tree.push_back(Node());
        int id = (int)tree.size() - 1;
        tree[id].value = value;
        tree[id].priority = (int)rng();
        tree[id].count = 1;
        tree[id].size = 1;
        return id;
    }

    int node_size(int u) const {
        return u == 0 ? 0 : tree[u].size;
    }

    void pushup(int u) {
        tree[u].size = node_size(tree[u].child[0]) +
                       node_size(tree[u].child[1]) +
                       tree[u].count;
    }

    // direction=0: lift left child by right rotation.
    // direction=1: lift right child by left rotation.
    void rotate(int &u, int direction) {
        int v = tree[u].child[direction];
        tree[u].child[direction] = tree[v].child[direction ^ 1];
        tree[v].child[direction ^ 1] = u;
        pushup(u);
        pushup(v);
        u = v;
    }

    void insert(int &u, int value) {
        if (u == 0) {
            u = new_node(value);
            return;
        }
        if (tree[u].value == value) {
            tree[u].count++;
            pushup(u);
            return;
        }

        int direction = value > tree[u].value;
        insert(tree[u].child[direction], value);
        if (tree[tree[u].child[direction]].priority < tree[u].priority) {
            rotate(u, direction);
        }
        pushup(u);
    }

    void insert(int value) {
        insert(root, value);
    }

    // Rank is 1-based: the smallest value has rank 1.
    int rank_of(int value) const {
        int u = root;
        int rank = 1;
        while (u != 0) {
            if (value <= tree[u].value) {
                u = tree[u].child[0];
            } else {
                rank += node_size(tree[u].child[0]) + tree[u].count;
                u = tree[u].child[1];
            }
        }
        return rank;
    }

    int kth(int k) const {
        int u = root;
        while (u != 0) {
            int left_size = node_size(tree[u].child[0]);
            if (k <= left_size) {
                u = tree[u].child[0];
            } else if (k <= left_size + tree[u].count) {
                return tree[u].value;
            } else {
                k -= left_size + tree[u].count;
                u = tree[u].child[1];
            }
        }
        return -1;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int m, n;
    cin >> m >> n;
    vector<int> a(m);
    for (int i = 0; i < m; i++) {
        cin >> a[i];
    }
    vector<int> u(n);
    for (int i = 0; i < n; i++) {
        cin >> u[i];
    }

    Treap treap(m + 5);
    int inserted = 0;
    int rank = 1; // 第 i 次 GET 输出第 i 小
    for (int q : u) {
        for (int i = inserted; i < q; i++) {
            treap.insert(a[i]);
        }
        inserted = q;
        cout << treap.kth(rank) << '\n';
        rank++;
    }

    return 0;
}

复杂度

每个元素移动常数次，总时间 $O((m+n)\log m)$ ，空间 $O(m)$ 。

平衡树解法插入和查询各 $O(\log m)$ ，总时间同为 $O((m+n)\log m)$ ，空间 $O(m)$ 。

总结

目标排名每次只增加 1，双堆分界也只需向右移动一个元素。

P1801 黑匣子

题目描述

Black Box 是一种原始的数据库。它可以储存一个整数数组，还有一个特别的变量 $i$ 。最开始的时候 Black Box 是空的．而 $i=0$ 。这个 Black Box 要处理一串命令。

命令只有两种：

ADD(x)：把 $x$ 元素放进 Black Box;
GET： $i$ 加 $1$ ，然后输出 Black Box 中第 $i$ 小的数。

记住：第 $i$ 小的数，就是 Black Box 里的数的按从小到大的顺序排序后的第 $i$ 个元素。

我们来演示一下一个有 $11$ 个命令的命令串。（如下表所示）

序号	操作	$i$	数据库	输出
1	`ADD(3)`	$0$	$3$	/
2	`GET`	$1$	$3$	$3$
3	`ADD(1)`	$1$	$1,3$	/
4	`GET`	$2$	$1,3$	$3$
5	`ADD(-4)`	$2$	$-4,1,3$	/
6	`ADD(2)`	$2$	$-4,1,2,3$	/
7	`ADD(8)`	$2$	$-4,1,2,3,8$	/
8	`ADD(-1000)`	$2$	$-1000,-4,1,2,3,8$	/
9	`GET`	$3$	$-1000,-4,1,2,3,8$	$1$
10	`GET`	$4$	$-1000,-4,1,2,3,8$	$2$
11	`ADD(2)`	$4$	$-1000,-4,1,2,2,3,8$	/

现在要求找出对于给定的命令串的最好的处理方法。ADD 命令共有 $m$ 个，GET 命令共有 $n$ 个。现在用两个整数数组来表示命令串：

$a_1,a_2,\cdots,a_m$ ：一串将要被放进 Black Box 的元素。例如上面的例子中 $a=[3,1,-4,2,8,-1000,2]$ 。
$u_1,u_2,\cdots,u_n$ ：表示第 $u_i$ 个元素被放进了 Black Box 里后就出现一个 GET 命令。例如上面的例子中 $u=[1,2,6,6]$ 。输入数据不用判错。

输入格式

第一行两个整数 $m$ 和 $n$ ，表示元素的个数和 GET 命令的个数。

第二行共 $m$ 个整数，从左至右第 $i$ 个整数为 $a_i$ ，用空格隔开。

第三行共 $n$ 个整数，从左至右第 $i$ 个整数为 $u_i$ ，用空格隔开。

输出格式

输出 Black Box 根据命令串所得出的输出串，一个数字一行。

输入输出样例 #1

输入 #1

text

1
2
37 4
3 1 -4 2 8 -1000 2
1 2 6 6

输出 #1

text

说明/提示

数据规模与约定

对于 $30\%$ 的数据， $1 \leq n,m \leq 10^{4}$ ；
对于 $50\%$ 的数据， $1 \leq n,m \leq 10^{5}$ ；
对于 $100\%$ 的数据， $1 \leq n,m \leq 2 \times 10^{5},|a_i| \leq 2 \times 10^{9}$ ，保证 $u$ 序列单调不降。

黑匣子

目录

题意

思路

Python 知识

代码

双堆解法（C++）

双堆解法（Python）

平衡树解法（C++）

复杂度

总结