#798894 (C++17) No.1796 木上のクーロン

提出ソース

結果

問題	No.1796 木上のクーロン
コンテスト
ユーザー	Nachia
提出日時	2022-09-18 21:01:34
言語	C++17 (gcc 13.3.0 + boost 1.89.0)
結果	WA
実行時間	-
コード長	9,232 bytes
記録記録タグの例: 初AC ショートコード純ショートコード純主流ショートコード最速実行時間
コンパイル時間	3,157 ms
コンパイル使用メモリ	118,600 KB
最終ジャッジ日時	2025-02-07 11:56:23
ジャッジサーバーID （参考情報）	judge5 / judge5

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

ファイルパターン	結果
sample	WA * 3
other	WA * 34

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

コンパイルメッセージ

main.cpp: In function ‘int main()’:
main.cpp:221:17: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  221 |     int N; scanf("%d", &N);
      |            ~~~~~^~~~~~~~~~
main.cpp:222:41: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  222 |     for(int i=0; i<N; i++){ int a; scanf("%d", &a); }
      |                                    ~~~~~^~~~~~~~~~
main.cpp:224:37: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  224 |     for(auto& e : E){ int u,v; scanf("%d%d", &u, &v); e = {u-1,v-1}; }
      |                                ~~~~~^~~~~~~~~~~~~~~~

ソースコード

raw source code


#include <utility>
#include <vector>
#include <algorithm>

namespace nachia{

template<class Elem>
class CsrArray{
public:
    struct ListRange{
        using iterator = typename std::vector<Elem>::iterator;
        iterator begi, endi;
        iterator begin() const { return begi; }
        iterator end() const { return endi; }
        int size() const { return (int)std::distance(begi, endi); }
        Elem& operator[](int i) const { return begi[i]; }
    };
    struct ConstListRange{
        using iterator = typename std::vector<Elem>::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 Elem& operator[](int i) const { return begi[i]; }
    };
private:
    int m_n;
    std::vector<Elem> m_list;
    std::vector<int> m_pos;
public:
    CsrArray() : m_n(0), m_list(), m_pos() {}
    static CsrArray Construct(int n, const std::vector<std::pair<int, Elem>>& items){
        CsrArray res;
        res.m_n = n;
        std::vector<int> buf(n+1, 0);
        for(auto& [u,v] : items){ ++buf[u]; }
        for(int i=1; i<=n; i++) buf[i] += buf[i-1];
        res.m_list.resize(buf[n]);
        for(int i=(int)items.size()-1; i>=0; i--){
            res.m_list[--buf[items[i].first]] = items[i].second;
        }
        res.m_pos = std::move(buf);
        return res;
    }
    static CsrArray FromRaw(std::vector<Elem> list, std::vector<int> pos){
        CsrArray res;
        res.m_n = pos.size() - 1;
        res.m_list = std::move(list);
        res.m_pos = std::move(pos);
        return res;
    }
    ListRange operator[](int u) { return ListRange{ m_list.begin() + m_pos[u], m_list.begin() + m_pos[u+1] }; }
    ConstListRange operator[](int u) const { return ConstListRange{ m_list.begin() + m_pos[u], m_list.begin() + m_pos[u+1] }; }
    int size() const { return m_n; }
    int fullSize() const { return (int)m_list.size(); }
};

} // namespace nachia
#include <cassert>
#include <tuple>

namespace nachia{

class StaticTopTree{
public:
    int m_n;
    nachia::CsrArray<int> m_children;
    std::vector<int> m_parent;
    std::vector<int> m_parentEdge;
    int m_root;

    struct Node{
        int p = -1, l = -1, r = -1;
        enum Type{ TyCompress, TyRake1, TyRake2, TyEdge } ty = TyEdge;
    };
    std::vector<Node> m_node;

    StaticTopTree(int n, std::vector<std::pair<int,int>> tree, int root = 0){
        m_n = n;
        assert((int)tree.size() == n-1);
        assert(0 <= root && root < n);
        m_root = root;
        nachia::CsrArray<int> adj; {
            std::vector<std::pair<int,int>> dedges(n*2-2);
            for(int i=0; i<n-1; i++){
                assert(0 <= tree[i].first && tree[i].first < n);
                assert(0 <= tree[i].second && tree[i].second < n);
                dedges[i*2] = tree[i];
                dedges[i*2+1] = { tree[i].second, tree[i].first };
            }
            adj = nachia::CsrArray<int>::Construct(n, dedges);
        }
        m_parent.assign(n, -2);
        m_parentEdge.assign(n, -1);
        std::vector<int> bfs = {m_root};
        bfs.reserve(n);
        m_parent[m_root] = -1;
        for(int i=0; i<(int)bfs.size(); i++){
            int p = bfs[i];
            for(int nx : adj[p]) if(m_parent[nx] == -2){
                m_parent[nx] = p;
                bfs.push_back(nx);
            }
        }
        for(int i=0; i<n; i++) assert(m_parent[i] != -2); // not connected
        for(int i=0; i<n-1; i++){
            auto& e = tree[i];
            if(m_parent[e.first] == e.second) std::swap(e.first, e.second);
            m_parentEdge[e.second] = i;
        }
        adj = nachia::CsrArray<int>::Construct(n, tree);
        std::vector<int> nd(n, 1);
        for(int i=n-1; i>=1; i--) nd[m_parent[bfs[i]]] += nd[bfs[i]];
        for(int p=0; p<n; p++) for(int e=1; e<adj[p].size(); e++) if(nd[adj[p][0]] < nd[adj[p][e]]) std::swap(adj[p][0], adj[p][e]);
        m_node.resize(n*2-3);
        for(int i=0; i<n-1; i++) m_node[i].ty = Node::TyEdge;
        std::vector<int> troot(n, -1);
        int trp = n*2-3;
        troot[bfs[0]] = --trp;
        for(int s : bfs) if(m_parent[s] < 0 || adj[m_parent[s]][0] != s){
            struct SzNode { int sz, vid, nx; };
            std::vector<SzNode> sznode;
            std::vector<int> Hid = {0};
            std::vector<int> boundarySize;
            if(m_parent[s] >= 0){
                sznode.push_back({ 1, m_parentEdge[s], adj[m_parent[s]][0] });
                Hid.push_back(Hid.back()+1);
            }
            for(int p=s; ; p=adj[p][0]){
                if(adj[p].size() == 0) break;
                for(int e=1; e<adj[p].size(); e++){
                    if(adj[adj[p][e]].size() == 0) sznode.push_back({ 1, m_parentEdge[adj[p][e]], adj[p][e] });
                    else sznode.push_back({ nd[adj[p][e]], -1, adj[p][e] });
                    Hid.push_back(Hid.back());
                }
                sznode.push_back({ 1, m_parentEdge[adj[p][0]], adj[p][0] });
                Hid.push_back(Hid.back() + 1);
            }
            boundarySize.assign(sznode.size()+1, 0);
            for(int i=0; i<(int)sznode.size(); i++) boundarySize[i+1] = boundarySize[i] + sznode[i].sz;
            struct QueNode{ int p, l, r; };
            std::vector<QueNode> Que = { { troot[s], 0, (int)sznode.size() } };
            Que.reserve(sznode.size() * 2);
            for(int i=0; i<(int)Que.size(); i++){
                int tp = Que[i].p, l = Que[i].l, r = Que[i].r;
                if(r-l == 1){
                    troot[sznode[l].nx] = tp;
                    continue;
                }
                int m = Que[i].l;
                while(boundarySize[m] - boundarySize[l] < boundarySize[r] - boundarySize[m+1]) m++;
                if(Hid[l] == Hid[m]) m_node[tp].ty = Node::TyRake2;
                else if(Hid[m] == Hid[r]) m_node[tp].ty = Node::TyRake1;
                else m_node[tp].ty = Node::TyCompress;
                int pl = (m-l == 1) ? sznode[l].vid : -1;
                if(pl == -1) pl = --trp;
                int pr = (r-m == 1) ? sznode[m].vid : -1;
                if(pr == -1) pr = --trp;
                m_node[tp].l = pl;
                m_node[tp].r = pr;
                m_node[pl].p = tp;
                m_node[pr].p = tp;
                Que.push_back({ pl, l, m });
                Que.push_back({ pr, m, r });
            }
        }
        std::swap(m_children, adj);
    }

    StaticTopTree() : StaticTopTree(2, {{0,1}}) {}
};

} // namespace nachia

#include <cstdio>
#include <array>
#include <atcoder/convolution>
#include <atcoder/modint>

using Modint = atcoder::static_modint<998244353>;

struct FPS{
    std::vector<Modint> A;
    int size() const { return A.size(); }

};

FPS operator+(const FPS& l, const FPS& r){
    FPS res;
    res.A.resize(std::max(l.size(), r.size()));
    for(int i=0; i<l.size(); i++) res.A[i] = l.A[i];
    for(int i=0; i<r.size(); i++) res.A[i] += r.A[i];
    return res;
}

FPS operator*(const FPS& l, const FPS& r){
    FPS res;
    res.A = atcoder::convolution(l.A, r.A);
    return res;
}

FPS operator<<(const FPS& l, int d){
    FPS res;
    res.A.resize(l.size() + d);
    for(int i=0; i<l.size(); i++) res.A[i+d] = l.A[i];
    return res;
}

using Val = std::array<std::array<FPS,2>,2>;

FPS getSolid(const Val& v){
    FPS res;
    res.A.resize(std::max({ v[0][0].size(), v[0][1].size()+1, v[1][0].size()+1, v[1][1].size()+2 }));
    for(int s : {0,1}) for(int t : {0,1}) for(int i=0; i<v[s][t].size(); i++) res.A[i+s+t] += v[s][t].A[i];
    return res;
}

int main(){
    int N; scanf("%d", &N);
    for(int i=0; i<N; i++){ int a; scanf("%d", &a); }
    std::vector<std::pair<int,int>> E(N-1);
    for(auto& e : E){ int u,v; scanf("%d%d", &u, &v); e = {u-1,v-1}; }
    auto tree = nachia::StaticTopTree(N, E);

    std::vector<Val> res(tree.m_node.size());

    using NodeTy = nachia::StaticTopTree::Node::Type;
    for(int i=0; i<(int)tree.m_node.size(); i++){
        auto& node = tree.m_node[i];
        if(node.ty == NodeTy::TyEdge){
            res[i][0][0] = FPS{ {1} };
            res[i][0][1] = FPS{ {1} };
            res[i][1][0] = FPS{ {1} };
            res[i][1][1] = FPS{ {} };
        }
        if(node.ty == NodeTy::TyCompress){
            for(int s : {0,1}) for(int t : {0,1}){
                res[i][s][t] = (res[node.l][s][0] * res[node.r][0][t]) + ((res[node.l][s][1] * res[node.r][1][t]) << 1);
            }
        }
        if(node.ty == NodeTy::TyRake1){
            for(int s : {0,1}) for(int t : {0,1}){
                res[i][s][t] = (res[node.l][s][t] * res[node.r][t][0]) + ((res[node.l][s][t] * res[node.r][t][1]) << 1);
            }
        }
        if(node.ty == NodeTy::TyRake2){
            for(int s : {0,1}) for(int t : {0,1}){
                res[i][s][t] = (res[node.l][s][0] * res[node.r][s][t]) + ((res[node.l][s][1] * res[node.r][s][t]) << 1);
            }
        }
    }

    FPS ans = getSolid(res.back());
    ans.A.resize(N+1);
    for(int i=0; i<=N; i++){
        if(i) printf(" ");
        printf("%u", ans.A[i].val());
    }
    printf("\n");
    return 0;
}

yukicoder

結果

コンパイルメッセージ

ソースコード