#1159427 (C++23) No.3506 All Distance is Square Number

• 問題: No.3506 All Distance is Square Number
• ユーザー: rumblycascade7
• 提出日時: 2026-04-18 11:18:37
• 言語: C++23 (gcc 15.2.0 + boost 1.89.0)
• 結果: AC
• 実行時間: 223 ms / 2,000 ms
• コード長: 19,621 bytes
• コンパイル時間: 3,282 ms
• コンパイル使用メモリ: 277,776 KB
• 実行使用メモリ: 42,028 KB
• 最終ジャッジ日時: 2026-04-18 11:18:54
• 合計ジャッジ時間: 7,347 ms
• ジャッジサーバーID （参考情報）: judge3_0 / judge1_0

ファイルパターン	結果
other	AC * 29

ソースコード

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <functional>
#include <iostream>
#include <stdexcept>
#include <unordered_map>
#include <utility>
#include <vector>

using namespace std;

struct Edge {
    int u;
    int v;
    int w;
};

static bool is_square(int x) {
    if (x < 0) return false;
    int r = static_cast<int>(sqrt(static_cast<long double>(x)));
    while ((r + 1) * 1LL * (r + 1) <= x) ++r;
    while (r * 1LL * r > x) --r;
    return r * 1LL * r == x;
}

static vector<uint16_t> normalize(vector<int> vals) {
    sort(vals.begin(), vals.end());
    vals.erase(unique(vals.begin(), vals.end()), vals.end());
    vector<uint16_t> out;
    out.reserve(vals.size());
    for (int x : vals) out.push_back(static_cast<uint16_t>(x));
    return out;
}

static bool contains_sum(const vector<uint16_t>& vals, int x) {
    if (x < 0) return false;
    return binary_search(vals.begin(), vals.end(), static_cast<uint16_t>(x));
}

static vector<int> reversed_path(const vector<int>& path) {
    return vector<int>(path.rbegin(), path.rend());
}

class Solver {
public:
    explicit Solver(int n) : N(n) {
        square.assign(11050, false);
        for (int k = 0; k * k < static_cast<int>(square.size()); ++k) square[k * k] = true;
        pair_paths.assign(N + 1, vector<vector<int>>(N + 1));
    }

    void solve() {
        if (N == 2) {
            solve_small_2();
            return;
        }
        if (N == 3) {
            solve_small_3();
            return;
        }
        if (N == 4) {
            solve_small_4();
            return;
        }
        if (N == 5) {
            solve_small_5();
            return;
        }

        init_base_edges();
        build_base_pair_paths();
        build_base_maps();
        init_stage_14();
        extend_edges_and_paths();
        print_answer();
    }

private:
    using SumList = vector<uint16_t>;
    using PathMap = unordered_map<int, vector<int>>;

    int N;
    vector<bool> square;
    vector<Edge> edges;
    vector<vector<vector<int>>> pair_paths;

    static constexpr int HELPER = 7;
    static constexpr int BASE_TAIL = 14;

    array<PathMap, BASE_TAIL + 1> H0{};
    array<PathMap, BASE_TAIL + 1> T0{};
    array<PathMap, BASE_TAIL + 1> N0{};

    vector<vector<SumList>> H;
    vector<vector<SumList>> T;
    vector<vector<SumList>> Nsum;
    vector<int> attach_a;
    vector<int> attach_b;
    vector<int> attach_edge_h;
    vector<int> attach_edge_t;

    unordered_map<long long, vector<int>> memoH;
    unordered_map<long long, vector<int>> memoT;
    unordered_map<long long, vector<int>> memoN;

    static long long make_key(int i, int t, int s) {
        return (static_cast<long long>(t) * 128LL + i) * 20000LL + s;
    }

    void solve_small_2() {
        edges = {{1, 2, 1}};
        pair_paths[1][2] = {1};
        print_answer();
    }

    void solve_small_3() {
        edges = {{1, 2, 16}, {2, 3, 9}, {1, 3, 25}};
        pair_paths[1][2] = {1};
        pair_paths[1][3] = {3};
        pair_paths[2][3] = {2};
        print_answer();
    }

    void solve_small_4() {
        edges = {
            {1, 3, 9},
            {1, 4, 27},
            {2, 3, 16},
            {2, 4, 36},
            {3, 4, 7},
        };
        pair_paths[1][2] = {1, 3};
        pair_paths[1][3] = {1};
        pair_paths[1][4] = {1, 5};
        pair_paths[2][3] = {3};
        pair_paths[2][4] = {4};
        pair_paths[3][4] = {1, 2};
        print_answer();
    }

    void solve_small_5() {
        edges = {
            {1, 2, 144},
            {1, 3, 45},
            {1, 4, 1},
            {2, 3, 81},
            {3, 5, 36},
            {4, 5, 64},
        };
        pair_paths[1][2] = {1};
        pair_paths[1][3] = {1, 4};
        pair_paths[1][4] = {3};
        pair_paths[1][5] = {2, 5};
        pair_paths[2][3] = {4};
        pair_paths[2][4] = {1, 2, 5, 6};
        pair_paths[2][5] = {1, 2, 5};
        pair_paths[3][4] = {5, 6};
        pair_paths[3][5] = {5};
        pair_paths[4][5] = {6};
        print_answer();
    }

    void init_base_edges() {
        static const Edge base[] = {
            {2, 3, 105},
            {1, 2, 37},
            {2, 4, 93},
            {1, 4, 68},
            {5, 6, 60},
            {4, 5, 132},
            {3, 5, 99},
            {3, 4, 27},
            {2, 6, 64},
            {1, 7, 8},
            {2, 7, 6},
            {1, 8, 2},
            {6, 8, 3},
            {1, 9, 5},
            {3, 9, 10},
            {3, 10, 7},
            {8, 10, 11},
            {7, 11, 12},
            {10, 11, 13},
            {9, 12, 14},
            {11, 12, 15},
            {5, 13, 17},
            {12, 13, 18},
            {6, 14, 19},
            {13, 14, 20},
        };
        edges.assign(base, base + 25);
    }

    vector<vector<pair<int, int>>> build_adj(int vertex_limit, int edge_limit) const {
        vector<vector<pair<int, int>>> adj(vertex_limit + 1);
        for (int id = 1; id <= edge_limit; ++id) {
            const Edge& e = edges[id - 1];
            if (e.u > vertex_limit || e.v > vertex_limit) continue;
            adj[e.u].push_back({e.v, id});
            adj[e.v].push_back({e.u, id});
        }
        return adj;
    }

    void build_base_pair_paths() {
        for (int n = 6; n <= min(N, BASE_TAIL); ++n) {
            int edge_limit = 2 * n - 3;
            auto adj = build_adj(n, edge_limit);
            for (int s = 1; s <= n; ++s) {
                for (int t = s + 1; t <= n; ++t) {
                    if (!pair_paths[s][t].empty()) continue;
                    vector<int> cur;
                    vector<int> answer;
                    vector<int> vis(n + 1, 0);
                    vis[s] = 1;
                    function<bool(int, int)> dfs = [&](int u, int sum) {
                        if (u == t) {
                            if (square[sum]) {
                                answer = cur;
                                return true;
                            }
                            return false;
                        }
                        for (auto [v, eid] : adj[u]) {
                            if (vis[v]) continue;
                            vis[v] = 1;
                            cur.push_back(eid);
                            if (dfs(v, sum + edges[eid - 1].w)) return true;
                            cur.pop_back();
                            vis[v] = 0;
                        }
                        return false;
                    };
                    if (!dfs(s, 0)) throw runtime_error("base pair path not found");
                    pair_paths[s][t] = answer;
                }
            }
        }
    }

    void enumerate_paths(
        int src,
        int dst,
        bool forbid_internal_helper,
        PathMap& out
    ) {
        auto adj = build_adj(BASE_TAIL, 25);
        vector<int> cur;
        vector<int> vis(BASE_TAIL + 1, 0);
        vis[src] = 1;
        function<void(int, int)> dfs = [&](int u, int sum) {
            if (u == dst) {
                out.emplace(sum, cur);
                return;
            }
            for (auto [v, eid] : adj[u]) {
                if (vis[v]) continue;
                if (forbid_internal_helper && v == HELPER && v != dst) continue;
                vis[v] = 1;
                cur.push_back(eid);
                dfs(v, sum + edges[eid - 1].w);
                cur.pop_back();
                vis[v] = 0;
            }
        };
        dfs(src, 0);
    }

    void build_base_maps() {
        for (int i = 1; i <= BASE_TAIL; ++i) {
            enumerate_paths(i, HELPER, false, H0[i]);
            enumerate_paths(i, BASE_TAIL, false, T0[i]);
            enumerate_paths(i, BASE_TAIL, true, N0[i]);
        }
    }

    void init_stage_14() {
        int stage_cap = max(N, BASE_TAIL);
        H.assign(stage_cap + 1, {});
        T.assign(stage_cap + 1, {});
        Nsum.assign(stage_cap + 1, {});
        attach_a.assign(stage_cap + 1, 0);
        attach_b.assign(stage_cap + 1, 0);
        attach_edge_h.assign(stage_cap + 1, 0);
        attach_edge_t.assign(stage_cap + 1, 0);

        H[BASE_TAIL].assign(BASE_TAIL + 1, {});
        T[BASE_TAIL].assign(BASE_TAIL + 1, {});
        Nsum[BASE_TAIL].assign(BASE_TAIL + 1, {});
        for (int i = 1; i <= BASE_TAIL; ++i) {
            vector<int> hs, ts, ns;
            hs.reserve(H0[i].size());
            ts.reserve(T0[i].size());
            ns.reserve(N0[i].size());
            for (const auto& [sum, _] : H0[i]) hs.push_back(sum);
            for (const auto& [sum, _] : T0[i]) ts.push_back(sum);
            for (const auto& [sum, _] : N0[i]) ns.push_back(sum);
            H[BASE_TAIL][i] = normalize(hs);
            T[BASE_TAIL][i] = normalize(ts);
            Nsum[BASE_TAIL][i] = normalize(ns);
        }
    }

    bool has_square_with_offset(const SumList& vals, int delta) const {
        for (uint16_t x : vals) {
            int s = static_cast<int>(x) + delta;
            if (s < static_cast<int>(square.size()) && square[s]) return true;
        }
        return false;
    }

    void extend_edges_and_paths() {
        vector<int> used(201, 0);
        for (const auto& e : edges) used[e.w] = 1;

        for (int t = BASE_TAIL; t < N; ++t) {
            vector<int> unused;
            for (int x = 1; x <= 200; ++x) {
                if (!used[x]) unused.push_back(x);
            }

            int choose_a = -1;
            int choose_b = -1;
            for (int ia = 0; ia < static_cast<int>(unused.size()) && choose_a == -1; ++ia) {
                int a = unused[ia];
                vector<char> ok_with_a(t + 1, 0);
                for (int i = 1; i <= t; ++i) {
                    ok_with_a[i] = has_square_with_offset(H[t][i], a);
                }
                for (int ib = ia + 1; ib < static_cast<int>(unused.size()); ++ib) {
                    int b = unused[ib];
                    bool good = true;
                    for (int i = 1; i <= t; ++i) {
                        if (ok_with_a[i]) continue;
                        if (!has_square_with_offset(T[t][i], b)) {
                            good = false;
                            break;
                        }
                    }
                    if (good) {
                        choose_a = a;
                        choose_b = b;
                        break;
                    }
                }
            }
            if (choose_a == -1) throw runtime_error("failed to extend graph");

            int v = t + 1;
            used[choose_a] = used[choose_b] = 1;
            attach_a[v] = choose_a;
            attach_b[v] = choose_b;
            attach_edge_h[v] = static_cast<int>(edges.size()) + 1;
            edges.push_back({HELPER, v, choose_a});
            attach_edge_t[v] = static_cast<int>(edges.size()) + 1;
            edges.push_back({t, v, choose_b});

            H[v].assign(v + 1, {});
            T[v].assign(v + 1, {});
            Nsum[v].assign(v + 1, {});

            for (int i = 1; i <= t; ++i) {
                if (i == HELPER) {
                    vector<int> ns;
                    ns.reserve(Nsum[t][i].size() + 1);
                    ns.push_back(choose_a);
                    for (uint16_t x : Nsum[t][i]) ns.push_back(static_cast<int>(x) + choose_b);
                    Nsum[v][i] = normalize(ns);
                    H[v][i] = {0};
                    T[v][i] = Nsum[v][i];
                } else {
                    vector<int> hs;
                    hs.reserve(H[t][i].size() + Nsum[t][i].size());
                    for (uint16_t x : H[t][i]) hs.push_back(x);
                    for (uint16_t x : Nsum[t][i]) hs.push_back(static_cast<int>(x) + choose_b + choose_a);
                    H[v][i] = normalize(hs);

                    vector<int> ns;
                    ns.reserve(Nsum[t][i].size());
                    for (uint16_t x : Nsum[t][i]) ns.push_back(static_cast<int>(x) + choose_b);
                    Nsum[v][i] = normalize(ns);

                    vector<int> ts;
                    ts.reserve(Nsum[v][i].size() + H[t][i].size());
                    for (uint16_t x : Nsum[v][i]) ts.push_back(x);
                    for (uint16_t x : H[t][i]) ts.push_back(static_cast<int>(x) + choose_a);
                    T[v][i] = normalize(ts);
                }
            }

            vector<int> hs_new;
            hs_new.reserve(Nsum[t][HELPER].size() + 1);
            hs_new.push_back(choose_a);
            for (uint16_t x : Nsum[t][HELPER]) hs_new.push_back(static_cast<int>(x) + choose_b);
            H[v][v] = normalize(hs_new);
            T[v][v] = {0};
            Nsum[v][v] = {0};

            for (int i = 1; i < v; ++i) {
                int best_sum = 1e9;
                int best_kind = -1;
                int best_prev_sum = -1;

                for (uint16_t x : H[t][i]) {
                    int total = static_cast<int>(x) + choose_a;
                    if (total < static_cast<int>(square.size()) && square[total]) {
                        if (total < best_sum || (total == best_sum && best_kind > 0)) {
                            best_sum = total;
                            best_kind = 0;
                            best_prev_sum = x;
                        }
                    }
                }
                for (uint16_t y : T[t][i]) {
                    int total = static_cast<int>(y) + choose_b;
                    if (total < static_cast<int>(square.size()) && square[total]) {
                        if (total < best_sum) {
                            best_sum = total;
                            best_kind = 1;
                            best_prev_sum = y;
                        }
                    }
                }

                if (best_kind == -1) throw runtime_error("failed to build pair path");
                if (best_kind == 0) {
                    auto path = build_H(i, t, best_prev_sum);
                    path.push_back(attach_edge_h[v]);
                    pair_paths[i][v] = move(path);
                } else {
                    auto path = build_T(i, t, best_prev_sum);
                    path.push_back(attach_edge_t[v]);
                    pair_paths[i][v] = move(path);
                }
            }
        }
    }

    vector<int> build_H(int i, int t, int sum) {
        long long key = make_key(i, t, sum);
        auto it = memoH.find(key);
        if (it != memoH.end()) return it->second;

        vector<int> res;
        if (t == BASE_TAIL) {
            auto base_it = H0[i].find(sum);
            if (base_it == H0[i].end()) throw runtime_error("missing H base path");
            res = base_it->second;
        } else {
            int prev = t - 1;
            int a = attach_a[t];
            int b = attach_b[t];
            int eh = attach_edge_h[t];
            int et = attach_edge_t[t];

            if (i == HELPER) {
                if (sum != 0) throw runtime_error("invalid H helper state");
            } else if (i == t) {
                if (sum == a) {
                    res = {eh};
                } else if (contains_sum(Nsum[prev][HELPER], sum - b)) {
                    res = {et};
                    auto tail = reversed_path(build_N(HELPER, prev, sum - b));
                    res.insert(res.end(), tail.begin(), tail.end());
                } else {
                    throw runtime_error("invalid H new-vertex state");
                }
            } else if (contains_sum(H[prev][i], sum)) {
                res = build_H(i, prev, sum);
            } else if (contains_sum(Nsum[prev][i], sum - b - a)) {
                res = build_N(i, prev, sum - b - a);
                res.push_back(et);
                res.push_back(eh);
            } else {
                throw runtime_error("invalid H recurrence");
            }
        }

        memoH.emplace(key, res);
        return res;
    }

    vector<int> build_N(int i, int t, int sum) {
        long long key = make_key(i, t, sum);
        auto it = memoN.find(key);
        if (it != memoN.end()) return it->second;

        vector<int> res;
        if (t == BASE_TAIL) {
            auto base_it = N0[i].find(sum);
            if (base_it == N0[i].end()) throw runtime_error("missing N base path");
            res = base_it->second;
        } else {
            int prev = t - 1;
            int a = attach_a[t];
            int b = attach_b[t];
            int eh = attach_edge_h[t];
            int et = attach_edge_t[t];

            if (i == t) {
                if (sum != 0) throw runtime_error("invalid N new-vertex state");
            } else if (i == HELPER) {
                if (sum == a) {
                    res = {eh};
                } else if (contains_sum(Nsum[prev][i], sum - b)) {
                    res = build_N(i, prev, sum - b);
                    res.push_back(et);
                } else {
                    throw runtime_error("invalid N helper state");
                }
            } else if (contains_sum(Nsum[prev][i], sum - b)) {
                res = build_N(i, prev, sum - b);
                res.push_back(et);
            } else {
                throw runtime_error("invalid N recurrence");
            }
        }

        memoN.emplace(key, res);
        return res;
    }

    vector<int> build_T(int i, int t, int sum) {
        long long key = make_key(i, t, sum);
        auto it = memoT.find(key);
        if (it != memoT.end()) return it->second;

        vector<int> res;
        if (t == BASE_TAIL) {
            auto base_it = T0[i].find(sum);
            if (base_it == T0[i].end()) throw runtime_error("missing T base path");
            res = base_it->second;
        } else {
            int prev = t - 1;
            int a = attach_a[t];
            int b = attach_b[t];
            int eh = attach_edge_h[t];
            int et = attach_edge_t[t];

            if (i == t) {
                if (sum != 0) throw runtime_error("invalid T new-vertex state");
            } else if (i == HELPER) {
                if (sum == a) {
                    res = {eh};
                } else if (contains_sum(Nsum[prev][i], sum - b)) {
                    res = build_N(i, prev, sum - b);
                    res.push_back(et);
                } else {
                    throw runtime_error("invalid T helper state");
                }
            } else if (contains_sum(H[prev][i], sum - a)) {
                res = build_H(i, prev, sum - a);
                res.push_back(eh);
            } else if (contains_sum(Nsum[prev][i], sum - b)) {
                res = build_N(i, prev, sum - b);
                res.push_back(et);
            } else {
                throw runtime_error("invalid T recurrence");
            }
        }

        memoT.emplace(key, res);
        return res;
    }

    void print_answer() const {
        int edge_count = static_cast<int>(edges.size());
        if (N >= 6 && N < BASE_TAIL) edge_count = 2 * N - 3;

        cout << edge_count << '\n';
        for (int i = 0; i < edge_count; ++i) {
            const auto& e = edges[i];
            cout << e.u << ' ' << e.v << ' ' << e.w << '\n';
        }
        for (int i = 1; i <= N; ++i) {
            for (int j = i + 1; j <= N; ++j) {
                const auto& p = pair_paths[i][j];
                cout << p.size();
                for (int eid : p) cout << ' ' << eid;
                cout << '\n';
            }
        }
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    Solver solver(N);
    solver.solve();
    return 0;
}