HDU 6059 Kanade's trio （字典树, 2017 Multi-Univ Training Contest 3)

Problem

含 N 个数字的 A 数组，求有多少个三元组 (i,j,k) 满足 $i\lt j\lt k$ 且 $(A_i \oplus A_j) < (A_j\oplus A_k)$

Limit

$1\le T \le 20$

$1\le \sum n\le 5\times 10^5$

$0\le A_i\le 2^{30}$

Idea

利用字典树维护前 k-1 个数。当前处理第 k 个数。

显然对于 k 与 i 的最高不相同位 $k_p$ 与 $i_p$ ：

当 $i_p=0$ , $k_p=1$ 时，该最高不相同位之前的 $i_{higher} = k_{higher}$ 。则 $j_{higher}$ 可以为任意数，均不对 i, k 更高位（指最高不相同位之前的高位，后同）的比较产生影响。而此时 $j_p$ 位必须为 0 才可保证不等式 $(A_i \oplus A_j) < (A_j\oplus A_k)$ 成立。

当 $i_p=1, k_p=0$ 时，$j_p$ 位必须为 1 ，更高位任意。

故利用数组 cnt[31][2] 统计每一位为 0 ，为 1 的有多少个（在前 K-1 个数中）。在字典树插入第 k 个数时，同时统计最高不相同位，即对于每次插入的 p 位为 num[p] (取值 0 或 1)，在同父节点对应的 1-num[p] 为根子树的所有节点均可作为 i 来寻找 j 以获取对答案的贡献。其中又仅要求 $j_p$ 与 $i_p$ ($i_p$ 值即 1-num[p]) 相同，故 $j_p$ 有 cnt[p][ 1-num[p] ] 种取值方案。

但是，同时需要注意 i 与 j 有在 A 数组的先后关系 ($i\lt j$) 需要保证。故在字典树中额外维护一个 Ext 点，记录将每次新加入的点与多少原有点可构成 i, j 关系。在后续计算贡献时去掉。

Code

#include<bits/stdc++.h>
using namespace std;
const int N = 500000 + 10;
const int Node_max = N * 31;
int T, n, num[30], a[N], cnt[31][2];
long long ext = 0, ans;
struct Node{
    int nxt[2];
    int cnt, ext;
} Trie[Node_max];
int Tsize;
void calc(int tmp, long long c) {
    ans += Trie[tmp].cnt * 1ll * (Trie[tmp].cnt - 1) / 2;
    ext += (c-Trie[tmp].cnt) * 1ll * Trie[tmp].cnt - Trie[tmp].ext;
}
void Trie_insert(int idx){
    int tmp = 0;
    for(int i=0;i<30;i++) {
        if(!Trie[tmp].nxt[ num[i] ]) {
            Trie[tmp].nxt[ num[i] ] = ++Tsize;
        }
        if(Trie[tmp].nxt[ 1-num[i] ]) {
            calc(Trie[tmp].nxt[ 1-num[i] ], cnt[i][ 1-num[i] ]);
        }
        tmp = Trie[tmp].nxt[ num[i] ];
        Trie[tmp].cnt++;
        Trie[tmp].ext += cnt[i][num[i]] - Trie[tmp].cnt;
    }
    return; 
}
int main()
{
    scanf("%d", &T);
    while(T-- && scanf("%d", &n)!=EOF)
    {
        memset(Trie, 0, Tsize * 16 + 16);
        memset(cnt, 0, sizeof(cnt));
        Tsize = 0;
        ans = 0;
        ext = 0;
        for(int i=1, tmp;i<=n;i++)
        {
            scanf("%d", &a[i]);
            tmp = a[i];
            for(int j=29;j>=0;j--) {
                num[j] = tmp%2;
                cnt[j][tmp%2]++;
                tmp /= 2;
            }
            Trie_insert(i);
        }
        printf("%lld\n", ans + ext);
    }
}