Cache Friendly Binary Search: cache 友好的二分搜索
前段时间面试,一面试官问我一个问题。写一个二分搜索。我一分钟以内就切完了,面试官让我在x86平台下将代码优化到极致. 我一下子就懵了。感觉多年的计算机白学了。
众所周知二分查找的主要瓶颈在 cache miss. 因为查找过程中第一次可能在 数组的 $\frac{1}{2}$ 处,第二次可能就跑到了数组的 $\frac{3}{4}$ 或者 $\frac{1}{4}$ 处,如果数组够长,那前面的几次查找几乎都会miss。
因此优化二分查找的效率,主要就在于优化 cache miss
我上网找了两天资料,确实找到了一种方法可以优化二分查找eytzinger binary search,可是限制条件是十分昂贵的:
- 需要对数组布局做一个重排(单次二分查找不会有任何优势)
- 只有在对数组做多次静态二分查找的情形下会带来cache优化。
我们先来详细看一下算法。再来仔细分析这个算法
eytzinger binary search
eytzinger memory layout
先来看一下eytzinger 内存布局。
先举一个例子,一个有序数组的[1,…,15] 所对应的完全二分查找树如下
如果我们将这颗树保存为 pointer-free 的完全二叉树的形式,可以得到下面的数组形式 记为 $b$, base-1,也就是
- $b[1]$ 为根节点
- $b[2k]$ 为 $b[k]$ 的左孩子,且 $b[2k] <= b[k]$
- $b[2k+1]$ 为 $b[k]$ 的右孩子,且 $b[2k +1] >=b[k]$
也就是说
对于一个有序数组的eytzinger 内存布局定义为它的完全二分搜索树的数组形式,类似表示在计算机世界中并不少见,比如 binary-heap就是这样表示的。
如何将有序数组改为eytzinger形式
其实不难发现 原数组是这颗完全二分搜索树的中序遍历,而eytzinger布局是这颗完全二分搜索树的层序遍历
const int n = 1e5;
int a[n], b[n+1];
int eytzinger(int i = 0, int k = 1) {
if (k <= n) {
i = eytzinger(i, 2 * k);
b[k] = a[i++];
i = eytzinger(i, 2 * k + 1);
}
return i;
}
在eytzinger 布局上的二分搜索实现
原理很简单了,
1. 从根节点开始,如果 b[1] >=target,转到left分支,否则转到right 分支
2. 直到到达一个非叶子节点
那么退出的时候一定对应着的是非叶子节点的 index,我们如何找到所对应的答案呢?
int search(int x) {
int k = 1;
while (k <= n) {
if (b[k] >= x)
k = 2 * k;
else
k = 2 * k + 1;
}
k >>= __builtin_ffs(~k);
return b[k];
}
其实很简单,我们来看跳到做分支的条件 $b[k]>=target$, 也就是说跳到的做分支的下标一定是满足答案的下标(不失一般性,这里我们的二分主要为cpp的lower_bound
)
因此我们只要找到最后一次跳到左分支的下标就行了, 注意跳到右分支一定会在末尾append 一个1,因此最后 $k$ 出循环的时候append了几个1就说明跳了几次右分支(记为 $s$),也就是对 $k$ 取反后的第一个为1的bit。我们对 $k$ 左移 $s$ 次找到它的祖先就好了。
branch-free 优化
while (k <= n)
k = 2 * k + (b[k] < x);
prefetch
while (k <= n) {
__builtin_prefetch(b + k * block_size);
k = 2 * k + (b[k] < x);
}
Memory allignment
alignas(64) int b[n+1];
code
#pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;
const int n = (1<<20);
const int block_size = 16; // = 64 / 4 = cache_line_size / sizeof(int)
alignas(64) int a[n], b[n+1];
int eytzinger(int i = 0, int k = 1) {
if (k <= n) {
i = eytzinger(i, 2 * k);
b[k] = a[i++];
i = eytzinger(i, 2 * k + 1);
}
return i;
}
int search(int x) {
int k = 1;
while (k <= n) {
__builtin_prefetch(b + k * block_size);
k = 2 * k + (b[k] < x);
}
k >>= __builtin_ffs(~k);
return k;
}
分析
整个算法有一个非常大的预处理开销($O(n)$),开销会比单次二分查找还要大
- 只有当对于一个 array 做多次静态 二分查询的时候才会有加速(多次的意思是 $q»n$)
- 算法对cache-line的大小和array size 比较敏感
reference
- eytzinger Binary Search
- Erik Demaine. Cache-Oblivious Algorithms and Data Structures
- Paul-Virak Khuong, Pat Morin. Array layouts for comparison-based searching
版权声明
本作品为作者原创文章,采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议
作者: taotao
仅允许非商业转载,转载请保留此版权声明,并注明出处