AtCoder Beginner Contest 252

Ex K-th Beautiful Necklace

题意

有 N 个石头，每个石头有不同的颜色和价值
颜色一共有 C 种，每种颜色至少存在一个石头
可以选择一些石头串成一条项链，项链的价值是所有石头价值的异或和
从 N 个石头中选择 C 个石头串一串项链，要求每个石头的颜色都不同
问所有的串法中，价值第 K 大的项链的价值是多少

分析

C 的最大值是70，假设每个颜色的石头都有2个或3个，那么可选的方案最大值为 \(2^{35}\) 或 \(3^{22} * 4\)
所以暴力枚举不可取
可以进行 meet-in-the-middle 处理，石头个数为 C 的项链可以由2串石头个数为 C/2 的项链拼接得到
假设前一串项链为 a，后一串项链为 b，从 a 中选择一串项链 x，求 b 中有多少串项链满足与 x 的异或大于某个值
01 字典树可以满足这个需求，对一个数字 x 求树中与 x 的异或值大于某个数 y 的数字的个数可以在 \( \log{y}\) 时间内求得
那么可以对答案进行二分，先用 b 构建字典树，然后对 a 中每个数 x 进行枚举，求字典树中大于 x 的数字的个数，总和与 K 相比
时间复杂度为 \( O(s\log^{2}V) \)，s 为 a 中项链的总数，V 为项链价值的最大值
这个时间复杂度无法接受
考虑另一种枚举方式，从答案的最高位开始枚举
考虑当前位为 1，然后枚举 a 的每个数字 x，可以在字典树中求出 考虑了前面所有位之后 与 x 异或后当前位为 1 的数字的个数，总和为 m
如果 m 大于等于 K，说明第 K 大的数字是在这 m 个数字中，所以答案的当前位应该取 1
如果 m 小于 K，说明第 K 大的数字不在这 m 个数字中，要更小，所以答案的当前位取 0
对于 x 来说，考虑到了第 i 位，那么就对应这个数字在字典树中的第 i-1 层到第 i 层的走向
所以当前位确定后需要对 x 的在字典树中的位置进行更新，根据答案的当前位和 x 的当前位决定沿字典树中0的边更新还是1的边更新
对于 m 小于 K 的情况，就是在剩下的数字中考虑第 K - m 大的数字，所以答案考虑当前位为1的时候需要 K -= m
总的时间复杂度为 \( s\log{V} \)

code

#include <bits/stdc++.h>
using namespace  std;
typedef long long ll;
typedef pair<int, int> pii;

const int N = 8e5 + 5;
const int mdep = 60;
int son[N*mdep][2];
int tot[N*mdep];
int nn = 1;
int cur[N];
vector<ll> color[75];
void insert(const ll x) {
    int cn = 1;
    tot[cn]++;
    for (int i = mdep;i >= 0;i--) {
        int cb = x >> i & 1LL;
        if (!son[cn][cb]) {
            son[cn][cb] = ++nn;
        }
        cn = son[cn][cb];
        tot[cn]++;
    }
}

int main() {
    int n,c;
    ll k;
    cin >> n >> c >> k;
    ll mul = 1;
    for (int i = 0;i < n;i++) {
        int col;
        ll val;
        cin >> col >> val;
        color[col].push_back(val);
    }
    for (int i = 1;i <= c;i++) {
        if (color[i].empty()) {
            continue;
        }
        mul *= (ll)color[i].size();
    }
    int max_col = 1;
    ll mul1 = 1;
    for (int i = 1;i <= c;i++) {
        if (!color[i].empty()) {
            mul1 *= color[i].size();
            if (mul / mul1 <= mul1) {
                max_col = i;
                break;
            }
        }
    }
    // meet in the middle
    vector<ll> tree_num;
    for (int i = max_col + 1;i <= c;i++) {
        if (color[i].empty()) {
            continue;
        }
        if (tree_num.empty()) {
            tree_num.swap(color[i]);
        } else {
            vector<ll> temp;
            for (const ll x : tree_num) {
                for (const ll y : color[i]) {
                    temp.push_back(x ^ y);
                }
            }
            tree_num.swap(temp);
        }
    }
    vector<ll> ntr_num;
    for (int i = 1;i <= max_col;i++) {
        if (color[i].empty()) {
            continue;
        }
        if (ntr_num.empty()) {
            ntr_num.swap(color[i]);
        } else {
            vector<ll> temp;
            for (const ll x : ntr_num) {
                for (const ll y : color[i]) {
                    temp.push_back(x ^ y);
                }
            }
            ntr_num.swap(temp);
        }
    }

    if (tree_num.empty()) {
        sort(ntr_num.begin(), ntr_num.end(), greater<>());
        cout << ntr_num[k-1] << endl;
    } else {
        ll ans = 0;
        // build 01-trie
        for (const ll x : tree_num) {
            insert(x);
        }
        for (int i = 0;i < ntr_num.size();i++) {
            cur[i] = 1;
        }
        for (int i = mdep;i >= 0;i--) {
            ll sum = 0;
            for (int j = 0;j < ntr_num.size();j++) {
                int cb = ntr_num[j] >> i & 1LL;
                int cn = cur[j];
                // find how many numbers xor is 1 for current bit position
                sum += tot[son[cn][cb^1]];
            }
            if (sum >= k) {
                ans |= (1LL << i);
                for (int j = 0;j < ntr_num.size();j++) {
                    // update x's position in 01-trie
                    int cb = ntr_num[j] >> i & 1LL;
                    int cn = cur[j];
                    cur[j] = son[cn][cb^1];
                }
            } else if (sum < k) {
                k -= sum;
                for (int j = 0;j < ntr_num.size();j++) {
                    // update x's position in 01-trie
                    int cb = ntr_num[j] >> i & 1LL;
                    int cn = cur[j];
                    cur[j] = son[cn][cb];
                }
            }
        }
        cout << ans << endl;
    }
    return 0;
}