前段时间面试,一面试官问我一个问题。写一个二分搜索。我一分钟以内就切完了,面试官让我在x86平台下将代码优化到极致. 我一下子就懵了。感觉多年的计算机白学了。

众所周知二分查找的主要瓶颈在 cache miss. 因为查找过程中第一次可能在 数组的 $\frac{1}{2}$ 处,第二次可能就跑到了数组的 $\frac{3}{4}$ 或者 $\frac{1}{4}$ 处,如果数组够长,那前面的几次查找几乎都会miss。

因此优化二分查找的效率,主要就在于优化 cache miss

我上网找了两天资料,确实找到了一种方法可以优化二分查找eytzinger binary search,可是限制条件是十分昂贵的:

  1. 需要对数组布局做一个重排(单次二分查找不会有任何优势)
  2. 只有在对数组做多次静态二分查找的情形下会带来cache优化。

我们先来详细看一下算法。再来仔细分析这个算法

eytzinger binary search

eytzinger memory layout

先来看一下eytzinger 内存布局。

先举一个例子,一个有序数组的[1,…,15] 所对应的完全二分查找树如下

如果我们将这颗树保存为 pointer-free 的完全二叉树的形式,可以得到下面的数组形式 记为 $b$, base-1,也就是

  1. $b[1]$ 为根节点
  2. $b[2k]$ 为 $b[k]$ 的左孩子,且 $b[2k] <= b[k]$
  3. $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)$),开销会比单次二分查找还要大

  1. 只有当对于一个 array 做多次静态 二分查询的时候才会有加速(多次的意思是 $q»n$)
  2. 算法对cache-line的大小和array size 比较敏感

reference

  1. eytzinger Binary Search
  2. Erik Demaine. Cache-Oblivious Algorithms and Data Structures
  3. Paul-Virak Khuong, Pat Morin. Array layouts for comparison-based searching

版权声明

本作品为作者原创文章,采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议

作者: taotao

仅允许非商业转载,转载请保留此版权声明,并注明出处