feature image

2022年3月3日 | ブログ記事

k-d tree のお手軽実装について

こんにちは、競技プログラミングをやっている tatyam と申します。私のお気に入りのデータ構造は -d tree です。

PAST-N : kDTree書くのたのし〜 pic.twitter.com/s5tB9e4piP

— tatyam (@tatyam_prime) May 3, 2020

kDTree さいこー!https://t.co/diMEs3KAhM

— tatyam (@tatyam_prime) March 28, 2021

E : kDTree ありがとう…
平面走査したくないので 3 次元 kDTree をする
O(n + m log m + k m^(2/3))
TLE する印象しかなかったんだけど 795 ms で通った

— tatyam (@tatyam_prime) June 7, 2021

できること

-d tree とは

-d tree のイメージ図。領域を縦横交互に、点の数が半分になるように分割する。

https://www.slideshare.net/okuraofvegetable/ss-65377588
https://ja.wikipedia.org/wiki/Kd木

とか読むと良いんじゃないでしょうか

お手軽 -d tree

点の挿入・削除がなく、矩形クエリしかしないときの -d tree を簡単に実装します。こういう木構造を簡単にするポイントは、葉木 (葉にしか要素を持たない木) にすることです。

void chmin(int& a, int b){ if(a > b) a = b; }
void chmax(int& a, int b){ if(a < b) a = b; }

struct kDTree{
    using T = pair<int, int>;
    using Iter = vector<T>::iterator;
    kDTree *l = nullptr, *r = nullptr;
    // 矩形クエリのために x / y 座標の最小値と最大値を持つ
    // 挿入・削除はないので、何の要素を持っているかは必要ない (要素を持っている葉と要素を持っていないそれ以外のノードを同じように扱える)
    int xmin = INT_MAX, xmax = INT_MIN, ymin = INT_MAX, ymax = INT_MIN, size = 0;
    kDTree(Iter begin, Iter end, bool divx = true){
        // vector<T> で渡してもいいですが、コピーコストがかかるので in-place に
        for(auto p = begin; p != end; p++){
            auto [x, y] = *p;
            chmin(xmin, x);
            chmax(xmax, x);
            chmin(ymin, y);
            chmax(ymax, y);
        }
        size = int(end - begin);
        if(size <= 1) return;
        // 葉木なので x / y 座標で半分に分けて再帰
        auto cen = begin + size / 2;
        // 縦横交互に分ける
        if(divx){
            nth_element(begin, cen, end, [](T a, T b){ return a.first < b.first; });
        }
        else{
            nth_element(begin, cen, end, [](T a, T b){ return a.second < b.second; });
        }
        l = new kDTree(begin, cen, !divx);
        r = new kDTree(cen, end, !divx);
    }
    // [x1, x2] * [y1, y2] にある点の個数を数える
    int count(int x1, int x2, int y1, int y2) const {
        // [xmin, xmax] * [ymin, ymax] と [x1, x2] * [y1, y2] に共通部分がない
        if(x2 < xmin || xmax < x1 || y2 < ymin || ymax < y1) return 0;
        // [xmin, xmax] * [ymin, ymax] 全体が [x1, x2] * [y1, y2] に含まれている
        if(x1 <= xmin && xmax <= x2 && y1 <= ymin && ymax <= y2) return size;
        // [xmin, xmax] * [ymin, ymax] の一部が [x1, x2] * [y1, y2] に含まれている -> 子に任せる
        return l->count(x1, x2, y1, y2) + r->count(x1, x2, y1, y2);
    }
};

追記

xmax - xminymax - ymin のうち長い方で分けるコードを最初は書いていましたが、このコードでは、 個の点が xmax - xmin ymax - ymin のような細長い領域にあると、取得クエリが かかってしまう場合があることを @noshi91 にご指摘いただきました。
しかし、これを想定して落とすことはかなり難しいので、AC してしまうことが多いと思われます。

問題を解く

PAST2 N - ビルの建設

次元平面上に正方形の領域が 個あり、 個目の領域は で、コストは です。
個のクエリに答えてください。

正方形領域に のコストを加算なので、imos 法を適用して累積和を取るだけにしましょう。
この後、疎な点集合の 次元累積和を取るのに、平面走査 + 座標圧縮 + BIT をしがちですが、Range Tree… は実装が大変ですし、その代わりに -d tree を使ってみましょう。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
void chmin(int& a, int b){ if(a > b) a = b; }
void chmax(int& a, int b){ if(a < b) a = b; }


using T = array<int, 3>;
struct kDTree{
    using Iter = vector<T>::iterator;
    kDTree *l = nullptr, *r = nullptr;
    int xmin = INT_MAX, xmax = INT_MIN, ymin = INT_MAX, ymax = INT_MIN;
    ll sum = 0;
    kDTree(Iter begin, Iter end, bool divx = true){
        for(auto p = begin; p != end; p++){
            auto [x, y, w] = *p;
            chmin(xmin, x);
            chmax(xmax, x);
            chmin(ymin, y);
            chmax(ymax, y);
            sum += w;
        }
        const int size = int(end - begin);
        if(size <= 1) return;
        auto cen = begin + size / 2;
        if(divx){
            nth_element(begin, cen, end, [](T& a, T& b){ return a[0] < b[0]; });
        }
        else{
            nth_element(begin, cen, end, [](T& a, T& b){ return a[1] < b[1]; });
        }
        l = new kDTree(begin, cen, !divx);
        r = new kDTree(cen, end, !divx);
    }
    // [-INF, x] * [-INF, y] にある点の重みを数える
    ll get(int x, int y) const {
        // [xmin, xmax] * [ymin, ymax] と [-INF, x] * [-INF, y] に共通部分がない
        if(x < xmin || y < ymin) return 0;
        // [xmin, xmax] * [ymin, ymax] 全体が [-INF, x] * [-INF, y] に含まれている
        if(xmax <= x && ymax <= y) return sum;
        // [xmin, xmax] * [ymin, ymax] の一部が [x1, x2] * [y1, y2] に含まれている -> 子に任せる
        return l->get(x, y) + r->get(x, y);
    }
};
int main(){
    int N, Q;
    cin >> N >> Q;
    vector<T> a;
    while(N--){
        int x, y, D, C;
        cin >> x >> y >> D >> C;
        D++;
        // imos 法
        a.push_back({x, y, C});
        a.push_back({x + D, y + D, C});
        a.push_back({x, y + D, -C});
        a.push_back({x + D, y, -C});
    }
    kDTree tree(a.begin(), a.end());
    while(Q--){
        int A, B;
        cin >> A >> B;
        cout << tree.get(A, B) << '\n';
    }
}

https://atcoder.jp/contests/past202004-open/submissions/29400829

その 2

-d tree は遅延セグ木のようなことができるということなので、矩形範囲にある点の重みを 加算するクエリで解いてみましょう。
加算クエリは可換なクエリなので、遅延評価部分を伝播させる必要がありません。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
void chmin(int& a, int b){ if(a > b) a = b; }
void chmax(int& a, int b){ if(a < b) a = b; }


using T = pair<int, int>;
struct kDTree{
    using Iter = vector<T>::iterator;
    kDTree *l = nullptr, *r = nullptr;
    int xmin = INT_MAX, xmax = INT_MIN, ymin = INT_MAX, ymax = INT_MIN;
    ll lazy = 0;
    kDTree(Iter begin, Iter end, bool divx = true){
        for(auto p = begin; p != end; p++){
            auto [x, y] = *p;
            chmin(xmin, x);
            chmax(xmax, x);
            chmin(ymin, y);
            chmax(ymax, y);
        }
        const int size = int(end - begin);
        if(size <= 1) return;
        auto cen = begin + size / 2;
        if(divx){
            nth_element(begin, cen, end, [](T& a, T& b){ return a.first < b.first; });
        }
        else{
            nth_element(begin, cen, end, [](T& a, T& b){ return a.second < b.second; });
        }
        l = new kDTree(begin, cen, !divx);
        r = new kDTree(cen, end, !divx);
    }
    // [x1, x2] * [y1, y2] にある点に C を加算
    void add(int x1, int x2, int y1, int y2, int C){
        // [xmin, xmax] * [ymin, ymax] と [x1, x2] * [y1, y2] に共通部分がない
        if(x2 < xmin || xmax < x1 || y2 < ymin || ymax < y1) return;
        // [xmin, xmax] * [ymin, ymax] 全体が [x1, x2] * [y1, y2] に含まれている
        if(x1 <= xmin && xmax <= x2 && y1 <= ymin && ymax <= y2){
            lazy += C;
            return;
        }
        // [xmin, xmax] * [ymin, ymax] の一部が [x1, x2] * [y1, y2] に含まれている
        l->add(x1, x2, y1, y2, C);
        r->add(x1, x2, y1, y2, C);
    }
    // [x, x] * [y, y] にある点の重みを数える
    ll get(int x, int y) const {
        // [xmin, xmax] * [ymin, ymax] と [x, x] * [y, y] に共通部分がない
        if(x < xmin || xmax < x || y < ymin || ymax < y) return 0;
        // [xmin, xmax] * [ymin, ymax] 全体が [x, x] * [y, y] に含まれている
        if(x == xmin && xmax == x && y == ymin && ymax == y) return lazy;
        // [xmin, xmax] * [ymin, ymax] が [x1, x2] * [y1, y2] を含んでいる
        return lazy + l->get(x, y) + r->get(x, y);
    }
};
int main(){
    int N, Q;
    cin >> N >> Q;
    vector<array<int, 4>> land(N);
    vector<T> query(Q);
    for(auto& [x, y, D, C] : land) cin >> x >> y >> D >> C;
    for(auto& [A, B] : query) cin >> A >> B;
    auto tree = [query]() mutable {
        return kDTree(query.begin(), query.end());
    }();
    for(auto [x, y, D, C] : land) tree.add(x, x + D, y, y + D, C);
    for(auto [A, B] : query) cout << tree.get(A, B) << '\n';
}

https://atcoder.jp/contests/past202004-open/submissions/29401260

ABC234 Ex - Enumerate Pairs

次元平面上に 個の点があります。ユークリッド距離が 以下であるような点の組を全て列挙してください。
(出力する組の個数)

ある点から距離 以下の点を列挙… これは -d tree !
本来は worst / query ですが、(出力する組の個数) の制約のおかげで、円形クエリであってもいい感じの計算量が保証されます。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
void chmin(int& a, int b){ if(a > b) a = b; }
void chmax(int& a, int b){ if(a < b) a = b; }


using T = array<int, 3>;
struct kDTree{
    using Iter = vector<T>::iterator;
    kDTree *l = nullptr, *r = nullptr;
    // 円形クエリのために x / y 座標の最小値と最大値を持つ
    int xmin = INT_MAX, xmax = INT_MIN, ymin = INT_MAX, ymax = INT_MIN;
    // 1 要素しか持っていない時の index
    int idx = -1;
    kDTree(Iter begin, Iter end){
        for(auto p = begin; p != end; p++){
            auto [x, y, i] = *p;
            chmin(xmin, x);
            chmax(xmax, x);
            chmin(ymin, y);
            chmax(ymax, y);
        }
        const int size = int(end - begin);
        if(size == 1) idx = (*begin)[2];
        if(size <= 1) return;
        auto cen = begin + size / 2;
        if((unsigned)xmax - (unsigned)xmin > (unsigned)ymax - (unsigned)ymin){
            // 長い方で分ける (正方形に近い方が円形クエリの効率がいい)
            nth_element(begin, cen, end, [](T& a, T& b){ return a[0] < b[0]; });
        }
        else{
            nth_element(begin, cen, end, [](T& a, T& b){ return a[1] < b[1]; });
        }
        l = new kDTree(begin, cen);
        r = new kDTree(cen, end);
    }
    // sqrt(dx^2 + dy^2) ≤ K
    static bool inside(int dx, int dy, int K){
        return ll(dx) * dx + ll(dy) * dy <= ll(K) * K;
    }
    // (x, y) から距離 K 以下の点を f に報告する
    template<class F> void get(int x, int y, int K, F f) const {
        // [xmin, xmax] * [ymin, ymax] と ((x, y) から半径 K 以内) に共通部分がない
        if(!inside(clamp(x, xmin, xmax) - x, clamp(y, ymin, ymax) - y, K)) return;
        // 葉なら idx を報告
        if(idx != -1){
            f(idx);
            return;
        }
        l->get(x, y, K, f);
        r->get(x, y, K, f);
    }
};
int main(){
    int N, K;
    cin >> N >> K;
    vector<T> A(N);
    for(auto& [x, y, i] : A) cin >> x >> y;
    for(int i = 0; i < N; i++) A[i][2] = i;
    vector<pair<int, int>> ans;
    kDTree tree(A.begin(), A.end());
    for(auto [x, y, i] : A) tree.get(x, y, K, [&, i = i](int j){ if(i < j) ans.emplace_back(i + 1, j + 1); });
    sort(ans.begin(), ans.end());
    cout << ans.size() << '\n';
    for(auto [x, y] : ans) cout << x << ' ' << y << '\n';
}

https://atcoder.jp/contests/abc234/submissions/29417326

tatyam icon
この記事を書いた人
tatyam

競技プログラミングと音ゲーをしています。 AtCoder銀冠 ICPC 2020/21 Yokohama 3 位, WF 15 位 (good_yamikin) ICPC 2021/22 Yokohama 3 位 (tonosama) ICPC 2022/23 Yokohama 1 位 (tonosama)

この記事をシェア

このエントリーをはてなブックマークに追加
共有

関連する記事

2023年7月13日
アルゴリズム班はやとき王選手権「競(けい)プロ」を開催しました!
abap34 icon abap34
2021年4月18日
ベズー係数とN項の拡張ユークリッドの互除法
0214sh7 icon 0214sh7
2023年4月29日
CPCTF2023 PPC作問陣 Writeup
noya2 icon noya2
2023年4月21日
CPCTFを開催します
noc7t icon noc7t
2022年8月30日
【競プロer向け】母関数を習得しよう!
tatyam icon tatyam
2021年4月26日
CPCTF2021 作問者writeup by hukuda222
hukuda222 icon hukuda222
記事一覧 タグ一覧 Google アナリティクスについて 特定商取引法に基づく表記