向量化线段树

masonxiong · 2026-10-21 16:20:31 · 算法·理论

省流：单点加，前缀和，树状数组两倍速度。

核心原理

将下标 i 拆成 w 进制 i=\sum_{k=0}^{h-1}i_k\cdot w^k，相当于把区间 [0,i) 拆成 i_k 个长度为 w^k 的区间。考虑一棵 w 叉线段树，自下向上第 k 层的节点管辖长度为 w^k\cdot w 的区间，节点上有一个长度为 w 的序列 s(k)，其中 s(k)_i 表示：这个长度为 w^k\cdot w 的区间中，前 i 个长度为 w^k 的小区间的所有元素之和。根据定义有 s(k)_0=0。

将第 k 层内所有节点的 s(k) 序列从左到右拼接，得到这一层的序列 \mathrm S(k)。它的下标是层内下标，不是原序列下标：原下标 i 落在 \mathrm S(k)_{\lfloor i\cdot w^{-k}\rfloor}。

设 j=\lfloor i\cdot w^{-k}\rfloor，则 \mathrm S(k)_j 是所有满足下列条件的 a_x 之和：把 x 与 i 写成 w 进制后，高于第 k 位的各位与 i 相同，第 k 位严格小于 i 的第 k 位，更低位任意。

下面是一个 w=4,n=16 的例子，最下一行为原序列 a：

[00,             10,             36,             78            ] S(1)
[00, 01, 03, 06, 00, 05, 11, 18, 00, 09, 19, 30, 00, 13, 27, 42] S(0)
[01, 02, 03, 04, 05, 06, 07, 08, 09, 10, 11, 12, 13, 14, 15, 16] a

查询 \sum[0,i) 就是求 \sum\limits_{x<i} a_x。每个 x<i 都有唯一的最高位 k，使得 x 与 i 在高于 k 的各位上相同，而 x 的第 k 位严格更小，更低位任意。

这些 x 的 a_x 之和就是 \mathrm S(k)_{\lfloor i\cdot w^{-k}\rfloor}。k 取遍 [0,h) 时，这些集合互不相交，并且覆盖全部 x<i，所以 [0,i) 的和为：

\sum_{k=0}^{h-1}\mathrm S(k)_{\lfloor i\cdot w^{-k}\rfloor}

例如 i=6=(12)_4，则 \sum[0,6)=\mathrm S(0)_6+\mathrm S(1)_{\lfloor 6\cdot 4^{-1}\rfloor}=11+10=21。于是求和的时间复杂度为 \mathcal O(h)=\mathcal O(\log_wn)。

接下来考虑修改 a_i\xleftarrow+d。第 k 层的 \mathrm S(k)_j 含有 a_i，当且仅当 j 与 \lfloor i\cdot w^{-k}\rfloor 位于同一个节点，且 j\bmod w>i_k。因此每一层只有一个节点会被改到：令 j=\lfloor i\cdot w^{-k}\rfloor，节点起点为 j-i_k，对节点内下标 v\in(i_k,w) 执行 \mathrm S(k)_{j-i_k+v}\xleftarrow+d 即可。于是修改的时间复杂度为 \mathcal O(wh)=\mathcal O(w\log_wn)。

注意到修改操作中每一层的修改形如"将连续 w 个元素进行增加"，这启发我们进行向量化。预处理掩码 m_{u,v}=-[u<v]，将 d 广播之后与 m_{i_k} 按位与再加到节点上即可。

接下来考虑 w 的选取。假设我们维护的是 32 位整数，使用 AVX2 指令集的 256 位向量，为了方便进行向量化，首先应当有 8\mid w。同时内存读取时总会以缓存行为最小单位进行取出，缓存行的大小为 64 字节，那我们刚好可以取 w=16，这样一个节点刚好为一个缓存行，同时修改操作可以拆成两次指令。

下面的代码给出一个示例，实现了 0-indexed 单点加，求左闭右开前缀和。冷知识：下面这个代码不依赖头文件。

:::info[示例代码]

template <int N>
class WideSGT {
    using Vector = int __attribute__((vector_size(32)));

    static constexpr int Cacheline = 64;
    static constexpr int Block = Cacheline / sizeof(int);
    static constexpr int Width = __builtin_ctz(Block);

    struct Mask {
        alignas(Cacheline) int mask[Block][Block];

        constexpr Mask() : mask{} {
            for (int u = 0; u < Block; ++u)
                for (int v = 0; v < Block; ++v)
                    mask[u][v] = u < v ? -1 : 0;
        }
    };

    static constexpr int height(int i) {
        int result = 1, n = i;
        while (n >= Block) {
            n >>= Width;
            ++result;
        }
        return result;
    }

    static constexpr int offset(int i) {
        int result = 0, n = N;
        while (i--) {
            n = (n + ~-Block) / Block;
            result += n * Block;
        }
        return result;
    }

    static constexpr int Size = N;
    static constexpr int Height = height(Size);

    alignas(Cacheline) int data[static_cast<unsigned>(offset(Height))];
    static constexpr Mask mask{};

public:
    int sum(int position) {
        int result = 0;
        for (int i = 0; i < Height; ++i) {
            result += data[offset(i) + position];
            position >>= Width;
        }
        return result;
    }

    void modify(int position, int delta) {
        auto vector = Vector{} + delta;
        for (int i = 0; i < Height; ++i) {
            auto block = (Vector*)&data[offset(i) + (position & -Block)];
            auto vmask = (Vector*)&mask.mask[position & ~-Block];
            block[0] += vector & vmask[0];
            block[1] += vector & vmask[1];
            position >>= Width;
        }
    }
};

:::

效率测试

评测机：Intel(R) Core(TM) Ultra 9 285H
编译器：g++ (x86_64-win32-seh-rev0, Built by MinGW-Builds project) 15.2.0
编译参数：-std=c++23 -march=native -mtune=native -Ofast

序列长度 n=10^7，2\cdot 10^7 次随机下标修改，2\cdot 10^7 次随机下标求和。10 次测试平均。

WideSGT：1162.4\space\rm ms
Fenwick：2052.3\space\rm ms

两倍速度。伟大无需多言。

:::info[测试代码]

template <int N>
class WideSGT {
    using Vector = int __attribute__((vector_size(32)));

    static constexpr int Cacheline = 64;
    static constexpr int Block = Cacheline / sizeof(int);
    static constexpr int Width = __builtin_ctz(Block);

    struct Mask {
        alignas(Cacheline) int mask[Block][Block];

        constexpr Mask() : mask{} {
            for (int u = 0; u < Block; ++u)
                for (int v = 0; v < Block; ++v)
                    mask[u][v] = u < v ? -1 : 0;
        }
    };

    static constexpr int height(int i) {
        int result = 1, n = i;
        while (n >= Block) {
            n >>= Width;
            ++result;
        }
        return result;
    }

    static constexpr int offset(int i) {
        int result = 0, n = N;
        while (i--) {
            n = (n + ~-Block) / Block;
            result += n * Block;
        }
        return result;
    }

    static constexpr int Size = N;
    static constexpr int Height = height(Size);

    alignas(Cacheline) int data[static_cast<unsigned>(offset(Height))];
    static constexpr Mask mask{};

public:
    int sum(int position) {
        int result = 0;
        for (int i = 0; i < Height; ++i) {
            result += data[offset(i) + position];
            position >>= Width;
        }
        return result;
    }

    void modify(int position, int delta) {
        auto vector = Vector{} + delta;
        for (int i = 0; i < Height; ++i) {
            auto block = (Vector*)&data[offset(i) + (position & -Block)];
            auto vmask = (Vector*)&mask.mask[position & ~-Block];
            block[0] += vector & vmask[0];
            block[1] += vector & vmask[1];
            position >>= Width;
        }
    }
};

template <int N>
class Fenwick {
    int data[static_cast<unsigned>(N)];

public:
    int sum(int position) {
        int result = 0;
        while (position--) {
            result += data[position];
            position &= -~position;
        }
        return result;
    }

    void modify(int position, int delta) {
        while (position < N) {
            data[position] += delta;
            position |= -~position;
        }
    }
};

#include <ctime>
#include <random>
#include <iostream>

auto main() -> int {
    static constexpr int N = 1E7;
    static constexpr int M = 2E7;
    static WideSGT<N> tree;

    int xorsum = 0;
    std::mt19937 rng;
    for (int i = M; i--; ) {
        tree.modify(rng() % N, 1);
        xorsum ^= i ^ tree.sum(rng() % N);
    }

    std::cerr << "XorSum = " << xorsum << '\n';
    std::cout << "Duration = " << std::clock() << " clocks\n";
}

:::