hdu 6121 Build a tree

Problem

acm.hdu.edu.cn/showproblem.php?pid=6121

Meaning

一棵 n 个点的完全 k 叉树，结点标号从 0 到 n - 1，求以每一棵子树的大小的异或和。

Analysis

一层层地统计答案。

找到一个特殊点：n - 1。以从它到树根的路径上的每一个点 p 为界，把树在该层的点劈成两半：左边是满 k 叉树的树根，右边是比左边的少了一层的满 k 叉树的树根，都很好算。
而 p 为根的子树，叶子层有可能是残缺的，除了叶子层外，也是一棵满 k 叉树。只要找到以 p 为根的子树中最左下角的后代结点编号，用 n - 1 一减，就得出子树 p 的叶子数。
要特判 k = 1 退化成链的情况（谜之打表找规律），照做应该会超时？
要注意各种爆 long long，比如：最后一层单独做、换乘法为除法…

Code

#include <cstdio>
using namespace std;
const int N = 1000000;

long long fast_pow(long long a, int b)
{
    long long res = 1;
    for( ; b; b >>= 1, a *= a)
        if(b & 1)
            res *= a;
    return res;
}

// 满 k 叉树，前 n 层总结点数
long long nNode(long long k, int n)
{
    return (fast_pow(k, n) - 1) / (k - 1);
}

long long table[100];

// table[i] = 满 k 叉树前 i 层总结点数
// 减少 nNode() 的调用（快速幂的次数）
// 虽然后来发现超时的原因并不在这
void da_biao(long long k, int d)
{
    for(int i = 0; i <= d; ++i)
        table[i] = nNode(k, i);
}

int main()
{
    int T;
    scanf("%d", &T);
    while(T--)
    {
        long long n, k, ans = 0;
        scanf("%I64d%I64d", &n, &k);
        // 特判 k = 1
        if(k == 1)
        {
            int rest = n & 3; // n % 4
            if(!rest)
                ans = n;
            else if(rest == 1)
                ans = 1;
            else if(rest == 2)
                ans = n + 1;
            else
                ans = 0;
            printf("%I64d\n", ans);
            continue;
        }
        // 求树的高度
        // 根的高度设为 1
        int depth = 1;
        for(long long m = n - 1; m > 0; ++depth)
            m = (m - 1) / k;

        da_biao(k, depth + 2);

        // 整棵树
        ans = n;
        // 最底层单独做
        ans ^= (n - table[depth - 1]) & 1;

        --depth;
        long long p = (n - 2) / k; // [(n - 1) - 1] / k，倒数第 2 层开始
        long long lid, rid, lnum, rnum, lch;
        for(int d = 2; p > 0; p = (p - 1) / k, ++d, --depth)
        {
            // 当前层最左边结点的标号
            lid = table[depth - 1];
            // 当前层最右边结点的标号
            rid = table[depth] - 1;
            // 左边的子树（满的）的大小
            lnum = table[d];
            // 右边的子树（少一层，但也是满的）的大小
            rnum = table[d - 1];

            if((p - lid) & 1)
                ans ^= lnum;
            if((rid - p) & 1)
                ans ^= rnum;

            lch = p; // p 为根的子数最左小角的后代结点
            while(lch <= (n - 2) / k) // lch * k + 1 <= n - 1
                lch = lch * k + 1;
            ans ^= table[d - 1] + n - lch;
        }
        printf("%I64d\n", ans);
    }
    return 0;
}