#905816 (C++14) No.2427 Tree Distance Two

提出ソース

結果

問題	No.2427 Tree Distance Two
コンテスト
ユーザー	Shirotsume
提出日時	2023-08-18 22:58:11
言語	C++14 (gcc 13.3.0 + boost 1.87.0)
結果	RE
実行時間	-
コード長	15,069 bytes
コンパイル時間	1,761 ms
コンパイル使用メモリ	112,076 KB
実行使用メモリ	17,332 KB
最終ジャッジ日時	2024-11-28 09:24:46
合計ジャッジ時間	5,928 ms
ジャッジサーバーID （参考情報）	judge5 / judge2

このコードへのチャレンジ
（要ログイン）

ファイルパターン	結果
other	WA * 20 RE * 15

権限があれば一括ダウンロードができます

コンパイルメッセージ

main.cpp: In constructor 'nachia::AdjacencyList::AdjacencyList(int, std::vector<std::pair<int, int> >, bool)':
main.cpp:25:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
   25 |         for(auto [u,v] : edges){ ++buf[u]; if(rev) ++buf[v]; }
      |                  ^
main.cpp:29:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
   29 |             auto [u,v] = edges[i];
      |                  ^
main.cpp: In constructor 'nachia::AdjacencyListEdgeIndexed::AdjacencyListEdgeIndexed(int, const std::vector<std::pair<int, int> >&, bool)':
main.cpp:88:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
   88 |         for(auto [u,v] : edges){ ++buf[u]; if(rev) ++buf[v]; }
      |                  ^
main.cpp:92:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
   92 |             auto [u,v] = edges[i];
      |                  ^
main.cpp: In member function 'nachia::AdjacencyListEdgeIndexed nachia::AdjacencyListEdgeIndexed::reversed_edges() const':
main.cpp:109:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
  109 |         for(auto [v,i] : E) ++buf[v];
      |                  ^
main.cpp:112:41: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17' [-Wc++17-extensions]
  112 |         for(int u=0; u<n; u++) for(auto [v,i] : operator[](u)) res.E[--buf[v]] = {u,i};
      |                                         ^

ソースコード

raw source code

#include <vector>
#include <utility>

namespace nachia{
    
struct AdjacencyList{
public:
    struct AdjacencyListRange{
        using iterator = typename std::vector<int>::const_iterator;
        iterator begi, endi;
        iterator begin() const { return begi; }
        iterator end() const { return endi; }
        int size() const { return (int)std::distance(begi, endi); }
        const int& operator[](int i) const { return begi[i]; }
    };
private:
    int mn;
    std::vector<int> E;
    std::vector<int> I;
public:
    AdjacencyList(int n, std::vector<std::pair<int,int>> edges, bool rev){
        mn = n;
        std::vector<int> buf(n+1, 0);
        for(auto [u,v] : edges){ ++buf[u]; if(rev) ++buf[v]; }
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        E.resize(buf[n]);
        for(int i=(int)edges.size()-1; i>=0; i--){
            auto [u,v] = edges[i];
            E[--buf[u]] = v;
            if(rev) E[--buf[v]] = u;
        }
        I = std::move(buf);
    }
    AdjacencyList(const std::vector<std::vector<int>>& edges = {}){
        int n = mn = edges.size();
        std::vector<int> buf(n+1, 0);
        for(int i=0; i<n; i++) buf[i+1] = buf[i] + edges[i].size();
        E.resize(buf[n]);
        for(int i=0; i<n; i++) for(int j=0; j<(int)edges[i].size(); j++) E[buf[i]+j] = edges[i][j];
        I = std::move(buf);
    }
    static AdjacencyList from_raw(std::vector<int> targets, std::vector<int> bounds){
        AdjacencyList res;
        res.mn = bounds.size() - 1;
        res.E = std::move(targets);
        res.I = std::move(bounds);
        return res;
    }
    AdjacencyListRange operator[](int u) const {
        return AdjacencyListRange{ E.begin() + I[u], E.begin() + I[u+1] };
    }
    int num_vertices() const { return mn; }
    int size() const { return num_vertices(); }
    int num_edges() const { return E.size(); }
    AdjacencyList reversed_edges() const {
        AdjacencyList res;
        int n = res.mn = mn;
        std::vector<int> buf(n+1, 0);
        for(int v : E) ++buf[v];
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        res.E.resize(buf[n]);
        for(int u=0; u<n; u++) for(int v : operator[](u)) res.E[--buf[v]] = u;
        res.I = std::move(buf);
        return res;
    }
};

struct AdjacencyListEdgeIndexed{
public:
    struct Edge { int to; int edgeidx; };
    struct AdjacencyListRange{
        using iterator = typename std::vector<Edge>::const_iterator;
        iterator begi, endi;
        iterator begin() const { return begi; }
        iterator end() const { return endi; }
        int size() const { return (int)std::distance(begi, endi); }
        const Edge& operator[](int i) const { return begi[i]; }
    };
private:
    int mn;
    std::vector<Edge> E;
    std::vector<int> I;
public:
    AdjacencyListEdgeIndexed(int n, const std::vector<std::pair<int,int>>& edges, bool rev){
        mn = n;
        std::vector<int> buf(n+1, 0);
        for(auto [u,v] : edges){ ++buf[u]; if(rev) ++buf[v]; }
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        E.resize(buf[n]);
        for(int i=(int)edges.size()-1; i>=0; i--){
            auto [u,v] = edges[i];
            E[--buf[u]] = { v, i };
            if(rev) E[--buf[v]] = { u, i };
        }
        I = std::move(buf);
    }
    AdjacencyListEdgeIndexed() : AdjacencyListEdgeIndexed(0, {}, false) {}
    AdjacencyListRange operator[](int u) const {
        return AdjacencyListRange{ E.begin() + I[u], E.begin() + I[u+1] };
    }
    int num_vertices() const { return mn; }
    int size() const { return num_vertices(); }
    int num_edges() const { return E.size(); }
    AdjacencyListEdgeIndexed reversed_edges() const {
        AdjacencyListEdgeIndexed res;
        int n = res.mn = mn;
        std::vector<int> buf(n+1, 0);
        for(auto [v,i] : E) ++buf[v];
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        res.E.resize(buf[n]);
        for(int u=0; u<n; u++) for(auto [v,i] : operator[](u)) res.E[--buf[v]] = {u,i};
        res.I = std::move(buf);
        return res;
    }
};

} // namespace nachia


namespace nachia{

template<unsigned int MOD>
struct PrimitiveRoot{
    static constexpr unsigned long long powm(unsigned long long a, unsigned long long i) {
        unsigned long long res = 1, aa = a;
        while(i){
            if(i & 1) res = res * aa % MOD;
            aa = aa * aa % MOD;
            i /= 2;
        }
        return res;
    }
    static constexpr bool examine_val(unsigned int g){
        unsigned int t = MOD - 1;
        for(unsigned long long d=2; d*d<=t; d++) if(t % d == 0){
            if(powm(g, (MOD - 1) / d) == 1) return false;
            while(t % d == 0) t /= d;
        }
        if(t != 1) if(powm(g, (MOD - 1) / t) == 1) return false;
        return true;
    }
    static constexpr unsigned int get_val(){
        for(unsigned int x=2; x<MOD; x++) if(examine_val(x)) return x;
        return 0;
    }
    static const unsigned int val = get_val();
};

}

#include <algorithm>



namespace nachia{

namespace freq_table_of_tree_dist_internal{

// a^i mod M
template<unsigned int MOD>
unsigned int powm(unsigned int a, unsigned long long i) {
    if(i == 0) return 1;
    unsigned int r = powm<MOD>((unsigned long long)a*a%MOD,i/2);
    if(i&1) r = (unsigned long long)r*a%MOD;
    return r;
}


template<unsigned int MOD, unsigned int g>
void NTT(std::vector<unsigned int>& A){
    int N = 1;
    while (N < (int)A.size()) N *= 2;
    static std::vector<unsigned int> exp_i_revbit_diff = { (unsigned int)powm<MOD>(g, (MOD - 1) >> 2) };
    for(int i=exp_i_revbit_diff.size(); (1<<(i+1))<N; i++){
        unsigned long long diffdiff = powm<MOD>(g, (MOD - 1) / 2 + ((MOD - 1) >> (i+2)) * 3);
        exp_i_revbit_diff.push_back(diffdiff);
    }
    for (int i = 1; i < N; i <<= 1) {
        unsigned long long q = 1;
        int maxk = N / i / 2;
        for (int j = 0; j < i; j++) {
            int off = j * maxk * 2;
            for (int k = off; k < off + maxk; k++) {
                unsigned int l = A[k];
                unsigned int r = A[k + maxk] * q % MOD;
                A[k] = std::min(l + r - MOD, l + r);
                A[k + maxk] = std::min(l - r, l + MOD - r);
            }
            if(j+1 != i) for(int d=0; ; d++) if(!(j&(1<<d))){
                q = q * exp_i_revbit_diff[d] % MOD;
                break;
            }
        }
    }
    for (int i = 0, j = 0; j < N; j++) {
        if (i < j) std::swap(A[i], A[j]);
        for (int k = N >> 1; k > (i ^= k); k >>= 1);
    }
}


template<unsigned int MOD>
std::vector<unsigned int> convolution_static(const std::vector<unsigned int>& A, const std::vector<unsigned int>& B){
    constexpr unsigned int g = nachia::PrimitiveRoot<MOD>::get_val();
    int Z = 1; while(Z < (int)(A.size() + B.size() - 1)) Z *= 2;
    std::vector<unsigned int> Ax(Z), Bx(Z);
    unsigned long long iZ = powm<MOD>(Z, MOD - 2);
    for(int i=0; i<(int)A.size(); i++) Ax[i] = A[i];
    for(int i=0; i<(int)B.size(); i++) Bx[i] = B[i];
    NTT<MOD, g>(Ax); NTT<MOD, g>(Bx);
    for(int i=0; i<Z; i++) Ax[i] = (unsigned long long)Ax[i] * Bx[i] % MOD;
    NTT<MOD, g>(Ax);
    reverse(Ax.begin() + 1, Ax.end());
    for(int i=0; i<Z; i++) Ax[i] = (unsigned long long)Ax[i] * iZ % MOD;
    Ax.resize(A.size() + B.size() - 1);
    return std::move(Ax);
}

static constexpr unsigned long long MOD0 = 998244353;
static constexpr unsigned long long MOD1 = 1107296257;

std::vector<unsigned long long> convolution(const std::vector<unsigned int>& A, const std::vector<unsigned int>& B){
    if(500 / A.size() >= B.size()){
        std::vector<unsigned long long> res(A.size() + B.size() - 1);
        for(std::size_t i=0; i<A.size(); i++) for(std::size_t j=0; j<B.size(); j++) res[i+j] += (unsigned long long)A[i] * B[j];
        return res;
    }
    std::vector<unsigned int> res0 = convolution_static<MOD0>(A,B);
    std::vector<unsigned int> res1 = convolution_static<MOD1>(A,B);
    std::vector<unsigned long long> res(res0.size());
    unsigned long long t = powm<MOD1>(MOD0, MOD1-2);
    for(int i=0; i<(int)res.size(); i++){
        unsigned long long x = res1[i] + MOD1 - res0[i];
        if(x >= MOD1) x -= MOD1;
        x = x * t % MOD1;
        res[i] = res0[i] + x * MOD0;
    }
  return res;
}



struct FrequencyTableOfTreeDistance{
    std::vector<unsigned long long> ans;
    std::vector<std::vector<unsigned int>> anstmp0;
    std::vector<std::vector<unsigned int>> anstmp1;
    static constexpr unsigned int NTTg0 = nachia::PrimitiveRoot<MOD0>::get_val();
    static constexpr unsigned int NTTg1 = nachia::PrimitiveRoot<MOD1>::get_val();

    void complete_ans(){
        unsigned long long t = powm<MOD1>(MOD0, MOD1-2);
        for(int i=0; i<(int)anstmp0.size(); i++){
            if(anstmp0[i].empty()) continue;
            int n = anstmp0[i].size();
            NTT<MOD0,NTTg0>(anstmp0[i]);
            std::reverse(anstmp0[i].begin()+1, anstmp0[i].end());
            NTT<MOD1,NTTg1>(anstmp1[i]);
            std::reverse(anstmp1[i].begin()+1, anstmp1[i].end());
            for(int j=0; j<std::min<int>(n, ans.size()); j++){
                unsigned long long x = anstmp1[i][j] + MOD1 - anstmp0[i][j];
                if(x >= MOD1) x -= MOD1;
                x = x * t % MOD1;
                ans[j] += anstmp0[i][j] + x * MOD0;
            }
            anstmp0[i].clear();
            anstmp1[i].clear();
        }
    }

    
    void append_convolute_once(std::vector<unsigned int> a, std::vector<unsigned int> b){
        int Z = 1, dZ = 0;
        while(Z < (int)(a.size() + b.size() - 1)){ Z *= 2; dZ++; }
        a.resize(Z, 0);
        b.resize(Z, 0);
        unsigned int invZ1 = powm<MOD1>(Z, MOD1-2);
        NTT<MOD1,NTTg1>(a);
        NTT<MOD1,NTTg1>(b);
        for(int i=0; i<(int)a.size(); i++) a[i] = (unsigned long long)a[i] * b[i] % MOD1 * invZ1 % MOD1;
        NTT<MOD1,NTTg1>(a);
        std::reverse(a.begin()+1, a.end());
        for(int i=0; i<(int)a.size(); i++) ans[i] += a[i];
    }


    void append_convolute(std::vector<unsigned int> a, std::vector<unsigned int> b){
        if((unsigned long long)a.size() * b.size() < 200){
            for(std::size_t i=0; i<a.size(); i++) for(std::size_t j=0; j<b.size(); j++){
                ans[i+j] += (unsigned long long)a[i] * b[j];
            }
            return;
        }
        int Z = 1, dZ = 0;
        while(Z < (int)(a.size() + b.size() - 1)){ Z *= 2; dZ++; }
        unsigned int invZ0 = powm<MOD0>(Z, MOD0-2);
        unsigned int invZ1 = powm<MOD1>(Z, MOD1-2);
        if(anstmp0[dZ].empty()){
            anstmp0[dZ].assign(Z, 0);
            anstmp1[dZ].assign(Z, 0);
        }
        a.resize(Z, 0);
        b.resize(Z, 0);
        std::vector<unsigned int> acpy = a;
        std::vector<unsigned int> bcpy = b;
        NTT<MOD0,NTTg0>(a);
        NTT<MOD0,NTTg0>(b);
        for(std::size_t i=0; i<a.size(); i++){
            unsigned long long tmp = ((unsigned long long)a[i] * b[i] % MOD0 * invZ0 + anstmp0[dZ][i]) % MOD0;
            anstmp0[dZ][i] = tmp;
        }
        a = std::move(acpy);
        b = std::move(bcpy);
        NTT<MOD1,NTTg1>(a);
        NTT<MOD1,NTTg1>(b);
        for(std::size_t i=0; i<a.size(); i++){
            unsigned long long tmp = ((unsigned long long)a[i] * b[i] % MOD1 * invZ1 + anstmp1[dZ][i]) % MOD1;
            anstmp1[dZ][i] = tmp;
        }
        anstmp1.push_back(std::move(a));
    }
    

    std::vector<unsigned long long> solve(const nachia::AdjacencyList& E){
        int n = E.num_vertices();
        int ceillog2n = 1;
        while(1 << ceillog2n < n) ceillog2n++;
        anstmp0.resize(ceillog2n + 2);
        anstmp1.resize(ceillog2n + 2);

        std::vector<int> Z(n,1);
        ans.resize(n,0);
        {
            std::vector<int> bfs = {0};
            bfs.reserve(n);
            std::vector<int> P(n, -1);
            for(int i=0; i<n; i++){
                int p = bfs[i];
                for(int e : E[p]) if(P[p] != e){
                    P[e] = p;
                    bfs.push_back(e);
                }
            }
            for(int i=n-1; i>=1; i--) Z[P[bfs[i]]] += Z[bfs[i]];
        }
        auto dfsmain = [&](int s, auto dfsmain)->void {
            if(Z[s] == 1){ Z[s] = 0; return; }
            while(true){
                int mxz = E[s].begin()[0];
                for(int e : E[s]) if(Z[e] > Z[mxz]) mxz = e;
                if(Z[mxz] * 2 > Z[s]){ Z[s] -= Z[mxz]; Z[mxz] += Z[s]; s = mxz; }
                else break;
            }
            Z[s] = 0;
            std::vector<std::vector<unsigned int>> dcnt;
            std::vector<int> dcnt_sum;
            for(int s2 : E[s]) if(Z[s2] != 0){
                std::vector<int> I = {s2};
                std::vector<int> D = {1};
                std::vector<int> P = {-1};
                for(int i=0; i<(int)I.size(); i++){
                    int p = I[i];
                    for(int e : E[p]) if(e != P[i]) if(Z[e] != 0){
                        I.push_back(e);
                        D.push_back(D[i]+1);
                        P.push_back(p);
                    }
                }
                std::vector<unsigned int> tmp(D.back()+1);
                for(int d : D) tmp[d]++;
                dcnt.emplace_back(std::move(tmp));
                dcnt_sum.emplace_back(Z[s2]);
            }

            std::sort(dcnt.begin(), dcnt.end(),
            [](const std::vector<unsigned int>& l, const std::vector<unsigned int>& r)->bool
                { return l.size() < r.size(); }
            );

            int dcnt_sum_sum = dcnt_sum[0];
            for(std::size_t i=0; i<dcnt.size()-1; i++){
                if(1'000'000'000 / dcnt_sum_sum > dcnt_sum[i+1]) append_convolute_once(dcnt[i+1], dcnt[i]);
                else append_convolute(dcnt[i+1], dcnt[i]);
                for(std::size_t k=0; k<dcnt[i].size(); k++) dcnt[i+1][k] += dcnt[i][k];
                dcnt_sum_sum += dcnt_sum[i+1];
            }
            for(std::size_t k=0; k<dcnt.back().size(); k++) ans[k] += dcnt.back()[k];
            
            for(int s2 : E[s]) if(Z[s2] != 0) dfsmain(s2,dfsmain);
        };
        dfsmain(0,dfsmain);

        complete_ans();
        return ans;
    }

};



} // freq_table_of_tree_dist_internal


std::vector<unsigned long long> FrequencyTableOfTreeDistance(const nachia::AdjacencyList& T){
    return freq_table_of_tree_dist_internal::FrequencyTableOfTreeDistance().solve(T);
}

} // namespace nachia

#include <iostream>
#define rep(i,n) for(int i=0; i<(n); i++)

int main(){
    using namespace std;
    int N; cin >> N;
    std::vector<pair<int,int>> E(N-1);
    rep(i,N-1){
        int u,v; cin >> u >> v;
        E[i] = make_pair(u,v);
    }
    auto ans = nachia::FrequencyTableOfTreeDistance(nachia::AdjacencyList(N,E,true));
    for(int i=1; i<=N-1; i++){
        if(i == 2) cout << ans[i];
    }
    cout << "\n";
    return 0;
}


struct ios_do_not_sync{
    ios_do_not_sync(){
        std::ios::sync_with_stdio(false);
        std::cin.tie(nullptr);
    }
} ios_do_not_sync_instance;

yukicoder

結果

コンパイルメッセージ

ソースコード