Manhattan MST
(lib/graph/manhattanMST.hpp)

• View this file on GitHub
• Last update: 2024-11-18 03:54:10+09:00
• Include: #include "lib/graph/manhattanMST.hpp"

Manhattan MST

概要

二次元平面上の $N$ 個 の頂点におけるマンハッタン距離での最小全域木を求めます。

(マンハッタン距離において最も近い 2 点を探したいときなどに使えます。)

計算量

• $\mathrm{O}(N \log N)$

Depends on

• Union-Find (lib/data_structure/union_find.hpp)

Verified with

• test/library_checker/geometry/manhattanmst.test.cpp

Code

#pragma once

#include "../data_structure/union_find.hpp"

/**
 * @brief Manhattan MST
 * @docs docs/graph/manhattanMST.md
 */

template <typename T>
vector<pair<int, int>> manhattanMST(vector<T> x, vector<T> y){
    int n = x.size();
    vector<tuple<T, int, int>> edge;
    edge.reserve(4 * n);
    vector<int> idx(n);
    iota(idx.begin(), idx.end(), 0);
    for(int s = 0; s < 2; ++s){
        for(int t = 0; t < 2; ++t){
            sort(idx.begin(), idx.end(), [&](const int i, const int j) {
                return x[i] + y[i] < x[j] + y[j];
            });
            map<T, int, greater<>> map;
            for(const int i : idx){
                for(auto iter = map.lower_bound(y[i]); iter != map.end(); iter = map.erase(iter)){
                    const int j = iter->second;
                    const T dx = x[i] - x[j];
                    const T dy = y[i] - y[j];
                    if(dy > dx) break;
                    edge.emplace_back(dx + dy, i, j);
                }
                map[y[i]] = i;
            }
            swap(x, y);
        }
        for(int i = 0; i < n; ++i) {
            x[i] *= -1;
        }
    }
    sort(edge.begin(), edge.end());
    UnionFind dsu(n);
    vector<pair<int, int>> used;
    used.reserve(n - 1);
    for(const auto &[c, i, j] : edge){
        if(!dsu.same(i, j)){
            used.emplace_back(i, j);
            dsu.unite(i, j);
        }
    }
    return used;
}

#line 2 "lib/graph/manhattanMST.hpp"

#line 2 "lib/data_structure/union_find.hpp"

/**
 * @brief Union-Find
 * @docs docs/data_structure/union_find.md
 */

#include <vector>
#include <cassert>

struct UnionFind{
    int V;
    std::vector<int> par;
    std::vector<int> edg;

    UnionFind(const int N) : V(N), par(N), edg(N){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            edg[i] = 0;
        }
    }

    int root(int x){
        assert(0 <= x && x < V);
        if(par[x] < 0) return x;
        return par[x] = root(par[x]);
    }

    int unite(int x, int y){
        int rx = root(x);
        int ry = root(y);
        if(rx == ry){
            edg[rx]++;
            return rx;
        }
        if(-par[rx] < -par[ry]) std::swap(rx, ry);
        par[rx] = par[rx] + par[ry];
        par[ry] = rx;
        edg[rx] += edg[ry] + 1;
        return rx;
    }

    bool same(int x, int y){
        int rx = root(x);
        int ry = root(y);
        return rx == ry;
    }

    long long size(int x){
        return -par[root(x)];
    }

    long long edge(int x){
        return edg[root(x)];
    }
};
#line 4 "lib/graph/manhattanMST.hpp"

/**
 * @brief Manhattan MST
 * @docs docs/graph/manhattanMST.md
 */

template <typename T>
vector<pair<int, int>> manhattanMST(vector<T> x, vector<T> y){
    int n = x.size();
    vector<tuple<T, int, int>> edge;
    edge.reserve(4 * n);
    vector<int> idx(n);
    iota(idx.begin(), idx.end(), 0);
    for(int s = 0; s < 2; ++s){
        for(int t = 0; t < 2; ++t){
            sort(idx.begin(), idx.end(), [&](const int i, const int j) {
                return x[i] + y[i] < x[j] + y[j];
            });
            map<T, int, greater<>> map;
            for(const int i : idx){
                for(auto iter = map.lower_bound(y[i]); iter != map.end(); iter = map.erase(iter)){
                    const int j = iter->second;
                    const T dx = x[i] - x[j];
                    const T dy = y[i] - y[j];
                    if(dy > dx) break;
                    edge.emplace_back(dx + dy, i, j);
                }
                map[y[i]] = i;
            }
            swap(x, y);
        }
        for(int i = 0; i < n; ++i) {
            x[i] *= -1;
        }
    }
    sort(edge.begin(), edge.end());
    UnionFind dsu(n);
    vector<pair<int, int>> used;
    used.reserve(n - 1);
    for(const auto &[c, i, j] : edge){
        if(!dsu.same(i, j)){
            used.emplace_back(i, j);
            dsu.unite(i, j);
        }
    }
    return used;
}

Back to top page