rANS：快速的渐进最优码

参考文献：

1. [Huff52] D. A. Huffman, “A Method for the Construction of Minimum-Redundancy Codes,” in Proceedings of the IRE, vol. 40, no. 9, pp. 1098-1101, Sept. 1952.
2. [Duda09] Jarek Duda. Asymmetric numeral systems. CoRR abs/0902.0271 (2009).
3. [Duda13] Jarek Duda. Asymmetric numeral systems: entropy coding combining speed of huffman coding with compression rate of arithmetic coding, 2013. ArXiv preprint, available at https://arxiv.org/abs/1311.2540.

为了压缩已知概率密度的符号序列，有两种主流方法：

1. 哈夫曼编码：使用二的幂次近似密度（二叉树），速度快，但压缩率低（除非信源密度恰好都是二的幂次）；
2. 算术/范围编码：使用精确密度，压缩率很容易达到香农熵，但是速度慢。

[Duda09] 提出了 Asymmetric numeral systems（ANS），它是渐进最优码（平均码长等于香农熵），并且速度比哈夫曼编码更快。

ANS

ANS 的抽象思路是：令 $[n]$ 是大小为 $n$ 的符号表，$q_s,s\in [n]$ 是其密度函数。对于一个充分大的整数 $x\in \mathbb{N}$，将其视为区间 x = { 0 , 1 , ⋯   , x − 1 } ⊆ N x=\{0,1,\cdots,x-1\} \subseteq \N x={0,1,⋯,x−1}⊆N，根据符号密度 $q_s$ 划分为 $n$ 个大小分别约为 $q_{s}x$ 的子集 $S_s\subseteq x$，满足 $x={\cup }_s{S}_s$，那么，在 $x$ 中均匀采样，就等价于首先根据 $q_s$ 采样出 $s$，然后在 $S_s$ 中均匀采样。所以，$x⟺(s,S_s)$ 存在双射，可以利用它来编码符号 $s$ 的序列。

函数 $D_1:\mathbb{N}\to [n],x\mapsto s$ 称为 distributing function，它给出了集合划分，
S s : = { y ∈ x ∣ D 1 ( y ) = s } S_s := \{y\in x \mid D_1(y)=s\} Ss​:={y∈x∣D1​(y)=s}
然后定义 $x_s:=|S_s|$ 以及 $D_2(x):=x_{D_1(x)}$

那么，解码函数为
$$D(x):=(D_1(x),D_2(x))=(s,x_s)$$
它是双射，编码函数是它的逆，
$$C(s,x_s):=x$$
为了实现渐进最优性，要求找到合适的函数 $D_1$，使得总有 $x\approx q_{s}x$

Range variants

[Duda09] 首先提出了一种 Uniform variants（uANS），函数 $D_1$ 将符号 $s$ 几乎均匀地铺在 $\mathbb{N}$ 上，其压缩率很好，但是计算相对较慢。下图展示了 $n=2,q_0=0.7,q_1=0.3$ 的情况，

之后，[Duda13] 提出了 rANS，函数 $D_1$ 将符号 $s$ 放置在一个区间内，其概率近似不如 uANS 精确，但是只需要更少的计算开销。

假设 $(g_0,g_1,\cdots ,g_{n-1})\in {\mathbb{N}}^n$ 是符号密度 $q_s$ 的良好近似（缩放 $B:={\sum }_s{g}_s$ 倍），将数组 { 0 , 1 , ⋯   , B − 1 } \{0,1,\cdots,B-1\} {0,1,⋯,B−1} 依次划分为大小为 $g_s$ 的区间。定义 $\text{CDF}(s):={\sum }_{i=0}^{s-1}{g}_i$ 是累积密度，初值 $\text{CDF}(0)=0$，那么范围 $[\text{CDF}(s),\text{CDF}(s+1))$ 被解码为 $s$

定义函数 $\text{symbol}:\mathbb{Z}\to [n]$，
$$\text{symbol}(x)=s⟺\text{CDF}(s)\le [x{]}_B<\text{CDF}(s+1)$$
编码函数设为：
$$C(s,x):=\lfloor \frac{x}{g_s}\rfloor \cdot B+[x{]}_{g_s}+\text{CDF}(s)\in \mathbb{N}$$
解码函数设为：
$$D(x):=\left(\text{symbol}(x),\lfloor \frac{x}{B}\rfloor \cdot g_s+[x{]}_B-\text{CDF}(s)\right)\in [n]\times \mathbb{N}$$
我在编程时注意到，如果初始 $x=0$ 并且序列开头有多个 $s=0$，那么总有 $C(s,x)=0$，因此解码时会出现 Bug；对于其他的 $x$ 值，是否存在类似问题？似乎设置 ${\sum }_s{g}_s\le B-1$ 以及 $CDF(0)=1$ 就可以解决该问题（但是这对么？）

假如选取 $B=2^t$，设置 $g_s\approx q_{s}B$，那么上述的 $x\cdot B,x/B,[x{]}_B$ 都可以用移位和掩码来快速实现。 然而 $x/g_s,x\cdot g_s,x+[x{]}_{g_s}$ 等运算，必须在大整数 $x$ 上处理，速度会相对较慢。

rANS 也是渐进最优码，满足：
$$\sum _{s\in [n]}{q}_s\log (C(s,x))\le \log x+\sum _{s\in [n]}{q}_s\log \frac{g_s}{B}+\frac{B}{x}$$
只要 $x$ 远大于 $B$，并且 $g_s/B$ 是 $q_s$ 的良好近似，那么平均码长就是香农熵。

Stream Encoding

对于无限符号序列，状态 $x\in \mathbb{N}$ 的比特长度也将无限增长，这使得计算效率降低。 [Duda09] 提出了流式编码，将状态限制在一个固定的区间内，从而提高计算效率，但是压缩率会略微变差。

我们将状态 $x$ 限制在范围 I : = { l , l + 1 , ⋯   , l b − 1 } I := \{l,l+1,\cdots,lb-1\} I:={l,l+1,⋯,lb−1}，易知 $|I|=(b-1)l$，

利用上述三条规则，

• 在编码时，
  1. 设置初值 $x=l\in I$，
  2. 根据 $C(s,x)$ 编码，直到即将 $x\ge lb$
  3. 不断将 $[x{]}_b$ 输出，并更新 $x\leftarrow \lfloor x/b\rfloor$，直到仍有 $x\in I$，然后回到 Step 2
  4. 在最后，将 $x\in I$ 输出
• 在解码时，
  1. 读取最近的状态 $x\in I$
  2. 根据 $D(x)$ 解码，直到即将 $x<l$
  3. 不断读取大小为 $b$ 的数据块 $d$，更新 $x\leftarrow xb+d$，直到仍有 $x\in I$，然后回到 Step 2
  4. 当 $x=l$ 且输入流为空（真的回到了初值），终止程序，并翻转解码结果

为了达到渐进最优性，需要让 $l$ 远大于 $B$。为了提高计算效率，可以将 $l,b$ 也都设置为二的幂次，尤其是 $w\mid \log b$（这里 $w=8/32$，取决于大数 $x\in \mathbb{N}$ 的基本运算单元）

代码实现

首先要自行实现一个大数类 INT，我这里使用数组 uint32_t *data 构造，

struct INT // 基底 2^32 的无符号大整数
{
    uint32_t *data = nullptr;
    int32_t len = 0;

    INT(int len); // 构造函数

    INT(const INT &x); // 深度复制

    ~INT(); // 析构函数

    int size(); // 当前大小（字节）

    void resize(int len); // 重新申请内存

    void clear(); // 清零

    void add(uint32_t x); // 加法

    void sub(uint32_t x); // 减法

    void mul(uint32_t x); // 乘法

    uint32_t div(uint32_t x); // 下取整除法，返回余数

    void shift(int i); // 按 word 左移（乘以 2^32），支持负数
};

主要的 rANS 实现如下：

void ANS_Build_CDF(uint32_t *CDF, const uint32_t *G, const int B)
{
    CDF[0] = 1; // 处理第一个符号是 s=0 的情况，否则 CDF[0]=0 总导致 x = 0，无法解码;
    for (int s = 1; s < B; s++)
        CDF[s] = CDF[s - 1] + G[s - 1]; // 构造累计密度分布表，CDF[s] = Pr[x < s]
}

uint32_t ANS_Symbol(uint32_t x, const uint32_t *CDF, const int B)
{
    int64_t L = 0, R = B;

    while (true) // 二分法查表，找出 CDF[M] <= x < CDF[M+1]
    {
        if (L >= B - 1)
            return B - 1;
        if (R <= 0)
            return 0;

        int64_t M = (L + R) / 2;
        if (x < CDF[M])
            R = M;
        else if (x >= CDF[M + 1])
            L = M + 1;
        else
            return M;
    }
}

void ANS_C(INT &x, const uint32_t s, const uint32_t *G, const uint32_t *CDF, const int B)
{
    uint32_t gs = G[s];
    uint32_t rem = x.div(gs); // 大数 x 除以 uint32 数据

    x.shift(1); // 乘以 B = 2^32，这里大数 x 的基本单元为 uint32
    x.add(rem);
    x.add(CDF[s]);
}

uint32_t ANS_D(INT &y, const uint32_t *G, const uint32_t *CDF, const int B)
{
    uint32_t lsb = y.data[0];
    uint32_t t = ANS_Symbol(lsb, CDF, B);

    y.shift(-1); // 整除 B，同理
    y.mul(G[t]);
    y.add(lsb);
    y.sub(CDF[t]);

    return t;
}

void ANS_Encode(vector<uint32_t> &x, const uint32_t *s, const int slen, const uint32_t *G, const uint32_t *CDF, const int B) // 原始版本，渐进最优码
{
    INT y(4);
    for (int i = 0; i < slen; i++)
        ANS_C(y, s[i], G, CDF, B); // 迭代编码符号序列

    int xlen = y.size();
    x.resize(xlen);
    memcpy(x.data(), y.data, xlen*sizeof(uint32_t));
}

void ANS_Decode(uint32_t **s, int &slen, const vector<uint32_t> &x, const uint32_t *G, const uint32_t *CDF, const int B)
{
    int xlen = x.size();
    INT y(x.size());
    memcpy(y.data, x.data(), xlen*sizeof(uint32_t));
    vector<uint32_t> res;

    while (y.size() != 0)
        res.push_back(ANS_D(y, G, CDF, B)); // 解出当前尾部符号

    if (*s != nullptr)
        free(*s);

    slen = res.size();
    *s = (uint32_t *)malloc(slen * sizeof(uint32_t));
    for (int i = 0; i < slen; i++)
        (*s)[i] = res[slen - i - 1]; // 翻转序列
}