ECR159-F Trees and XOR Queries Again 【xor基底】

問題

https://codeforces.com/contest/1902/problem/F (2400diff)

 N頂点からなる木があって、各頂点 vには整数 a_{v}が書き込まれている。

 Q個のクエリ (x, y, k)に答えよ

  • 木の x-yパス上の頂点(端点も含む)に書き込まれた整数からなる集合を Xとする。 Xの部分集合 X'であって、 X'の各要素の総xorが kになるようなものが存在するか判定せよ

制約:

  •  2\ \le\ N, Q\ \le\ 2\times 10^{5}
  •  0\ \le\ A_{v}\ \lt\ 2^{20}
  • TL6.5秒

解法

「xor基底とはなにか」、「xor基底の構築方法」、「xor基底の性質」について、公式解説に初心者向けの説明があった。ありがたい

頂点 1を根とする。各頂点 vについて vから根までのパス上の頂点の値からxor基底を構築することを考える。以下の疑似コードにしたがってxor基底を構築する

std::array<int, 20> buildBase(int v) {
   std::array<int, 20> B;
   B.fill(0);
   int x = v;
   while (x != -1) {
      if (a[x] is not in span(B)) {
           update B; // span(B)にa[x]も含まれるように更新する。やり方は公式解説を参照せよ。
      }
      x = parent of x;
   }
   return B;
}

各頂点 vについて、buildBase(v)を呼ぶ。普通にやると O(N^2)かかり死亡。

各頂点 vについて、buildBase(v)内で「update B」を呼び出すような頂点 xの列を保持する。これをgood(v)と呼ぶことにする。実はbuildBaseは以下のように改良することができる

std::array<int, 20> buildBase2(int v) {
   std::array<int, 20> B;
   B.fill(0);
   insert v to good[v];
   update B; // span(B)にa[v]を含まれるようにする。
   if (v is root) return B;
   for (vertex x : good(parent of v)) {
      if (a[x] is not in span(B)) {
           update B;
           insert x to good[v];
      }
   }
   return B;
}

buildBase(v)を実行したときに、update Bを呼び出さなかった x vの子孫 cbuildBase(c)を呼び出しても xでupdate Bされないことから言える。

基底の要素数は高々 \log (\max a)個(本問題では 20)であることから、各頂点について根から頂点までのパス上の値からなる基底を O(\log^{2} (\max a))で構築することができる。

クエリ処理を考える。good(x)good(y)に含まれる頂点から基底を構築して、 kが含まれるか確認すれば良い。

  • 当然だが、 \text{lca}(x, y)の子孫でない頂点は無視する

各クエリ O(\log^2(\max a) + \log N)で処理できる。

実装

#include <bits/stdc++.h>

#include "Src/Template/IOSetting.hpp"
using namespace zawa;

using base = std::array<int, 20>;

int highbit(int v) {
    for (int i{19} ; i >= 0 ; i--) if (v & (1 << i)) {
        return i;
    }
    return -1;
}

int reduce(const base& b, int v) {
    for (int i{19} ; i >= 0 ; i--) if (v & (1 << i)) {
        v ^= b[i];
    }
    return v;
}

bool insert(base& b, int v) {
    v = reduce(b, v);
    if (v == 0) return false;
    int bit{highbit(v)};
    assert(bit != -1);
    assert(b[bit] == 0);
    b[bit] = v;
    return true;
}

int main() {
    SetFastIO();

    int n; std::cin >> n;
    std::vector<int> a(n);
    for (auto& x : a) std::cin >> x;
    std::vector<std::vector<int>> g(n);
    for (int _{} ; _ < n - 1 ; _++) {
        int u, v; std::cin >> u >> v;
        u--; v--;
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }
    std::vector<int> par(n), in(n), out(n), dep(n);
    int timer{};
    auto dfs{[&](auto dfs, int v, int p) -> void {
        par[v] = (p == -1 ? v : p);
        in[v] = timer++;
        dep[v] = (p == -1 ? 0 : dep[p] + 1);
        for (auto x : g[v]) if (x != p) {
            dfs(dfs, x, v);
        }
        out[v] = timer;
    }};
    dfs(dfs, 0, -1);
    std::vector dob(19, std::vector<int>(n));
    for (int i{} ; i < n ; i++) {
        dob[0][i] = par[i];
    }
    for (int i{1} ; i < 19 ; i++) {
        for (int v{} ; v < n ; v++) {
            dob[i][v] = dob[i - 1][dob[i - 1][v]];
        }
    }
    std::vector<base> B(n);
    std::vector<std::vector<int>> good(n);
    for (auto& x : B) x.fill(0);
    auto dfs2{[&](auto dfs, int v, int p) -> void {
        if (a[v]) assert(insert(B[v], a[v]));
        good[v].emplace_back(v);
        if (p != -1) {
            for (auto x : good[p]) {
                if (insert(B[v], a[x])) {
                    good[v].emplace_back(x);
                }
            }
        }
        for (auto x : g[v]) if (x != p) dfs(dfs, x, v);
    }};
    dfs2(dfs2, 0, -1);
    auto lca{[&](int u, int v) -> int {
        if (dep[u] > dep[v]) std::swap(u, v);
        for (int i{18} ; i >= 0 ; i--) if ((dep[v] - dep[u]) & (1 << i)) {
            v = dob[i][v];
        }
        if (u == v) return u;
        for (int i{18} ; i >= 0 ; i--) if (dob[i][u] != dob[i][v]) {
            u = dob[i][u];
            v = dob[i][v];
        }
        assert(u != v);
        assert(dob[0][u] == dob[0][v]);
        return dob[0][u];
    }};
    auto isAnc{[&](int u, int v) -> bool {// vはuの祖先?
        return in[v] <= in[u] and out[u] <= out[v];
    }};
    int q; std::cin >> q;
    while (q--) {
        int u, v, x; std::cin >> u >> v >> x;
        if (x == 0) {
            std::cout << "YES" << '\n';
            continue;
        }
        u--; v--;
        int l{lca(u, v)};
        base b;
        b.fill(0);
        for (auto y : good[u]) if (isAnc(y, l)) {
            insert(b, a[y]);
        }
        for (auto y : good[v]) if (isAnc(y, l)) {
            insert(b, a[y]);
        }
        x = reduce(b, x);
        std::cout << (x ? "NO" : "YES") << '\n';
    }        
}

感想

バチャ中はhttps://atcoder.jp/contests/abc223/tasks/abc223_hの解説とにらめっこしていたが、解けず....

  • 列だからうまくいく解法であって、木の場合に拡張できなさそう
  • 反対に本問題はbit数が小さいことに注目した解法なので、こっちに適用すると破滅しそう

他のxor基底を使う問題の解説でxor基底の構築方法に言及しているものを見たことがなかったので、本問題の公式解説は非常に良い収穫だった。

実装時バグ取りが 2分で終わった。この実装の上手さを本番で出せれば良いんだけどな。