#897227 (C++23) No.2471 Gemini Tree(Ver.Lapislazuli)

提出ソース
結果

問題	No.2471 Gemini Tree(Ver.Lapislazuli)
ユーザー	👑 rin204
提出日時	2023-08-01 21:33:43
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	73 ms / 2,000 ms
コード長	18,153 bytes
コンパイル時間	3,918 ms
コンパイル使用メモリ	264,300 KB
実行使用メモリ	17,920 KB
最終ジャッジ日時	2024-11-08 06:07:14
合計ジャッジ時間	6,362 ms
ジャッジサーバーID （参考情報）	judge1 / judge4
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 5,248 KB
testcase_01	AC	2 ms 5,248 KB
testcase_02	AC	2 ms 5,248 KB
testcase_03	AC	2 ms 5,248 KB
testcase_04	AC	2 ms 5,248 KB
testcase_05	AC	2 ms 5,248 KB
testcase_06	AC	3 ms 5,248 KB
testcase_07	AC	2 ms 5,248 KB
testcase_08	AC	2 ms 5,248 KB
testcase_09	AC	2 ms 5,248 KB
testcase_10	AC	3 ms 5,248 KB
testcase_11	AC	3 ms 5,248 KB
testcase_12	AC	3 ms 5,248 KB
testcase_13	AC	3 ms 5,248 KB
testcase_14	AC	2 ms 5,248 KB
testcase_15	AC	3 ms 5,248 KB
testcase_16	AC	2 ms 5,248 KB
testcase_17	AC	2 ms 5,248 KB
testcase_18	AC	2 ms 5,248 KB
testcase_19	AC	2 ms 5,248 KB
testcase_20	AC	2 ms 5,248 KB
testcase_21	AC	2 ms 5,248 KB
testcase_22	AC	2 ms 5,248 KB
testcase_23	AC	52 ms 11,008 KB
testcase_24	AC	51 ms 11,136 KB
testcase_25	AC	53 ms 11,008 KB
testcase_26	AC	50 ms 11,136 KB
testcase_27	AC	53 ms 11,136 KB
testcase_28	AC	51 ms 11,008 KB
testcase_29	AC	54 ms 11,136 KB
testcase_30	AC	52 ms 11,136 KB
testcase_31	AC	53 ms 11,136 KB
testcase_32	AC	73 ms 17,920 KB
testcase_33	AC	49 ms 14,592 KB
testcase_34	AC	36 ms 12,160 KB
testcase_35	AC	38 ms 12,672 KB
testcase_36	AC	41 ms 13,440 KB
testcase_37	AC	55 ms 16,256 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
// start A.cpp
// #pragma GCC target("avx2")
// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>
using namespace std;

using ll  = long long;
using ull = unsigned long long;
template <class T>
using pq = priority_queue<T>;
template <class T>
using qp = priority_queue<T, vector<T>, greater<T>>;
#define vec(T, A, ...) vector<T> A(__VA_ARGS__);
#define vvec(T, A, h, ...) vector<vector<T>> A(h, vector<T>(__VA_ARGS__));
#define vvvec(T, A, h1, h2, ...) vector<vector<vector<T>>> A(h1, vector<vector<T>>(h2, vector<T>(__VA_ARGS__)));

#ifndef RIN__LOCAL
#define endl "\n"
#endif
#define spa ' '
#define len(A) A.size()
#define all(A) begin(A), end(A)

#define fori1(a) for (ll _ = 0; _ < (a); _++)
#define fori2(i, a) for (ll i = 0; i < (a); i++)
#define fori3(i, a, b) for (ll i = (a); i < (b); i++)
#define fori4(i, a, b, c) for (ll i = (a); ((c) > 0 || i > (b)) && ((c) < 0 || i < (b)); i += (c))
#define overload4(a, b, c, d, e, ...) e
#define fori(...) overload4(__VA_ARGS__, fori4, fori3, fori2, fori1)(__VA_ARGS__)

vector<char> stoc(string &S) {
    int n = S.size();
    vector<char> ret(n);
    for (int i = 0; i < n; i++) ret[i] = S[i];
    return ret;
}

#define INT(...)                                                                                                                                                                                       \
    int __VA_ARGS__;                                                                                                                                                                                   \
    inp(__VA_ARGS__);
#define LL(...)                                                                                                                                                                                        \
    ll __VA_ARGS__;                                                                                                                                                                                    \
    inp(__VA_ARGS__);
#define STRING(...)                                                                                                                                                                                    \
    string __VA_ARGS__;                                                                                                                                                                                \
    inp(__VA_ARGS__);
#define CHAR(...)                                                                                                                                                                                      \
    char __VA_ARGS__;                                                                                                                                                                                  \
    inp(__VA_ARGS__);
#define VEC(T, A, n)                                                                                                                                                                                   \
    vector<T> A(n);                                                                                                                                                                                    \
    inp(A);
#define VVEC(T, A, n, m)                                                                                                                                                                               \
    vector<vector<T>> A(n, vector<T>(m));                                                                                                                                                              \
    inp(A);

const ll MOD1 = 1000000007;
const ll MOD9 = 998244353;

template <class T>
auto min(const T &a) {
    return *min_element(all(a));
}
template <class T>
auto max(const T &a) {
    return *max_element(all(a));
}
template <class T, class S>
auto clamp(T &a, const S &l, const S &r) {
    return (a > r ? r : a < l ? l : a);
}
template <class T, class S>
inline bool chmax(T &a, const S &b) {
    return (a < b ? a = b, 1 : 0);
}
template <class T, class S>
inline bool chmin(T &a, const S &b) {
    return (a > b ? a = b, 1 : 0);
}
template <class T, class S>
inline bool chclamp(T &a, const S &l, const S &r) {
    auto b = clamp(a, l, r);
    return (a != b ? a = b, 1 : 0);
}

void FLUSH() {
    cout << flush;
}
void print() {
    cout << endl;
}
template <class Head, class... Tail>
void print(Head &&head, Tail &&...tail) {
    cout << head;
    if (sizeof...(Tail)) cout << spa;
    print(forward<Tail>(tail)...);
}
template <typename T>
void print(vector<T> &A) {
    int n = A.size();
    for (int i = 0; i < n; i++) {
        cout << A[i];
        if (i != n - 1) cout << ' ';
    }
    cout << endl;
}
template <typename T>
void print(vector<vector<T>> &A) {
    for (auto &row : A) print(row);
}
template <typename T, typename S>
void print(pair<T, S> &A) {
    cout << A.first << spa << A.second << endl;
}
template <typename T, typename S>
void print(vector<pair<T, S>> &A) {
    for (auto &row : A) print(row);
}
template <typename T, typename S>
void prisep(vector<T> &A, S sep) {
    int n = A.size();
    for (int i = 0; i < n; i++) {
        cout << A[i];
        if (i != n - 1) cout << sep;
    }
    cout << endl;
}
template <typename T, typename S>
void priend(T A, S end) {
    cout << A << end;
}
template <typename T>
void priend(T A) {
    priend(A, spa);
}
template <typename T, typename S>
bool printif(bool f, T A, S B) {
    if (f)
        print(A);
    else
        print(B);
    return f;
}
template <class... T>
void inp(T &...a) {
    (cin >> ... >> a);
}
template <typename T>
void inp(vector<T> &A) {
    for (auto &a : A) cin >> a;
}
template <typename T>
void inp(vector<vector<T>> &A) {
    for (auto &row : A) inp(row);
}
template <typename T, typename S>
void inp(pair<T, S> &A) {
    inp(A.first, A.second);
}
template <typename T, typename S>
void inp(vector<pair<T, S>> &A) {
    for (auto &row : A) inp(row.first, row.second);
}

template <typename T>
T sum(vector<T> &A) {
    T tot = 0;
    for (auto a : A) tot += a;
    return tot;
}

template <typename T>
vector<T> compression(vector<T> X) {
    sort(all(X));
    X.erase(unique(all(X)), X.end());
    return X;
}

vector<vector<int>> read_edges(int n, int m, bool direct = false, int indexed = 1) {
    vector<vector<int>> edges(n, vector<int>());
    for (int i = 0; i < m; i++) {
        INT(u, v);
        u -= indexed;
        v -= indexed;
        edges[u].push_back(v);
        if (!direct) edges[v].push_back(u);
    }
    return edges;
}
vector<vector<int>> read_tree(int n, int indexed = 1) {
    return read_edges(n, n - 1, false, indexed);
}
template <typename T>
vector<vector<pair<int, T>>> read_wedges(int n, int m, bool direct = false, int indexed = 1) {
    vector<vector<pair<int, T>>> edges(n, vector<pair<int, T>>());
    for (int i = 0; i < m; i++) {
        INT(u, v);
        T w;
        inp(w);
        u -= indexed;
        v -= indexed;
        edges[u].push_back({v, w});
        if (!direct) edges[v].push_back({u, w});
    }
    return edges;
}
template <typename T>
vector<vector<pair<int, T>>> read_wtree(int n, int indexed = 1) {
    return read_wedges<T>(n, n - 1, false, indexed);
}

inline bool yes(bool f = true) {
    cout << (f ? "yes" : "no") << endl;
    return f;
}
inline bool Yes(bool f = true) {
    cout << (f ? "Yes" : "No") << endl;
    return f;
}
inline bool YES(bool f = true) {
    cout << (f ? "YES" : "NO") << endl;
    return f;
}

inline bool no(bool f = true) {
    cout << (!f ? "yes" : "no") << endl;
    return f;
}
inline bool No(bool f = true) {
    cout << (!f ? "Yes" : "No") << endl;
    return f;
}
inline bool NO(bool f = true) {
    cout << (!f ? "YES" : "NO") << endl;
    return f;
}

// start other/Modint.hpp

template <int MOD>
struct Modint {
    int x;
    Modint() : x(0) {}
    Modint(int64_t y) {
        if (y >= 0)
            x = y % MOD;
        else
            x = (y % MOD + MOD) % MOD;
    }

    Modint &operator+=(const Modint &p) {
        x += p.x;
        if (x >= MOD) x -= MOD;
        return *this;
    }

    Modint &operator-=(const Modint &p) {
        x -= p.x;
        if (x < 0) x += MOD;
        return *this;
    }

    Modint &operator*=(const Modint &p) {
        x = int(1LL * x * p.x % MOD);
        return *this;
    }

    Modint &operator/=(const Modint &p) {
        *this *= p.inverse();
        return *this;
    }

    Modint &operator%=(const Modint &p) {
        assert(p.x == 0);
        return *this;
    }

    Modint operator-() const {
        return Modint(-x);
    }

    Modint &operator++() {
        x++;
        if (x == MOD) x = 0;
        return *this;
    }

    Modint &operator--() {
        if (x == 0) x = MOD;
        x--;
        return *this;
    }

    Modint operator++(int) {
        Modint result = *this;
        ++*this;
        return result;
    }

    Modint operator--(int) {
        Modint result = *this;
        --*this;
        return result;
    }

    friend Modint operator+(const Modint &lhs, const Modint &rhs) {
        return Modint(lhs) += rhs;
    }

    friend Modint operator-(const Modint &lhs, const Modint &rhs) {
        return Modint(lhs) -= rhs;
    }

    friend Modint operator*(const Modint &lhs, const Modint &rhs) {
        return Modint(lhs) *= rhs;
    }

    friend Modint operator/(const Modint &lhs, const Modint &rhs) {
        return Modint(lhs) /= rhs;
    }

    friend Modint operator%(const Modint &lhs, const Modint &rhs) {
        assert(rhs.x == 0);
        return Modint(lhs);
    }

    bool operator==(const Modint &p) const {
        return x == p.x;
    }

    bool operator!=(const Modint &p) const {
        return x != p.x;
    }

    bool operator<(const Modint &rhs) const {
        return x < rhs.x;
    }

    bool operator<=(const Modint &rhs) const {
        return x <= rhs.x;
    }

    bool operator>(const Modint &rhs) const {
        return x > rhs.x;
    }

    bool operator>=(const Modint &rhs) const {
        return x >= rhs.x;
    }

    Modint inverse() const {
        int a = x, b = MOD, u = 1, v = 0, t;
        while (b > 0) {
            t = a / b;
            a -= t * b;
            u -= t * v;
            swap(a, b);
            swap(u, v);
        }
        return Modint(u);
    }

    Modint pow(int64_t k) const {
        Modint ret(1);
        Modint y(x);
        while (k > 0) {
            if (k & 1) ret *= y;
            y *= y;
            k >>= 1;
        }
        return ret;
    }

    friend ostream &operator<<(ostream &os, const Modint &p) {
        return os << p.x;
    }

    friend istream &operator>>(istream &is, Modint &p) {
        int64_t y;
        is >> y;
        p = Modint<MOD>(y);
        return (is);
    }

    static int get_mod() {
        return MOD;
    }
};

struct Arbitrary_Modint {
    int x;
    static int MOD;

    static void set_mod(int mod) {
        MOD = mod;
    }

    Arbitrary_Modint() : x(0) {}
    Arbitrary_Modint(int64_t y) {
        if (y >= 0)
            x = y % MOD;
        else
            x = (y % MOD + MOD) % MOD;
    }

    Arbitrary_Modint &operator+=(const Arbitrary_Modint &p) {
        x += p.x;
        if (x >= MOD) x -= MOD;
        return *this;
    }

    Arbitrary_Modint &operator-=(const Arbitrary_Modint &p) {
        x -= p.x;
        if (x < 0) x += MOD;
        return *this;
    }

    Arbitrary_Modint &operator*=(const Arbitrary_Modint &p) {
        x = int(1LL * x * p.x % MOD);
        return *this;
    }

    Arbitrary_Modint &operator/=(const Arbitrary_Modint &p) {
        *this *= p.inverse();
        return *this;
    }

    Arbitrary_Modint &operator%=(const Arbitrary_Modint &p) {
        assert(p.x == 0);
        return *this;
    }

    Arbitrary_Modint operator-() const {
        return Arbitrary_Modint(-x);
    }

    Arbitrary_Modint &operator++() {
        x++;
        if (x == MOD) x = 0;
        return *this;
    }

    Arbitrary_Modint &operator--() {
        if (x == 0) x = MOD;
        x--;
        return *this;
    }

    Arbitrary_Modint operator++(int) {
        Arbitrary_Modint result = *this;
        ++*this;
        return result;
    }

    Arbitrary_Modint operator--(int) {
        Arbitrary_Modint result = *this;
        --*this;
        return result;
    }

    friend Arbitrary_Modint operator+(const Arbitrary_Modint &lhs, const Arbitrary_Modint &rhs) {
        return Arbitrary_Modint(lhs) += rhs;
    }

    friend Arbitrary_Modint operator-(const Arbitrary_Modint &lhs, const Arbitrary_Modint &rhs) {
        return Arbitrary_Modint(lhs) -= rhs;
    }

    friend Arbitrary_Modint operator*(const Arbitrary_Modint &lhs, const Arbitrary_Modint &rhs) {
        return Arbitrary_Modint(lhs) *= rhs;
    }

    friend Arbitrary_Modint operator/(const Arbitrary_Modint &lhs, const Arbitrary_Modint &rhs) {
        return Arbitrary_Modint(lhs) /= rhs;
    }

    friend Arbitrary_Modint operator%(const Arbitrary_Modint &lhs, const Arbitrary_Modint &rhs) {
        assert(rhs.x == 0);
        return Arbitrary_Modint(lhs);
    }

    bool operator==(const Arbitrary_Modint &p) const {
        return x == p.x;
    }

    bool operator!=(const Arbitrary_Modint &p) const {
        return x != p.x;
    }

    bool operator<(const Arbitrary_Modint &rhs) {
        return x < rhs.x;
    }

    bool operator<=(const Arbitrary_Modint &rhs) {
        return x <= rhs.x;
    }

    bool operator>(const Arbitrary_Modint &rhs) {
        return x > rhs.x;
    }

    bool operator>=(const Arbitrary_Modint &rhs) {
        return x >= rhs.x;
    }

    Arbitrary_Modint inverse() const {
        int a = x, b = MOD, u = 1, v = 0, t;
        while (b > 0) {
            t = a / b;
            a -= t * b;
            u -= t * v;
            swap(a, b);
            swap(u, v);
        }
        return Arbitrary_Modint(u);
    }

    Arbitrary_Modint pow(int64_t k) const {
        Arbitrary_Modint ret(1);
        Arbitrary_Modint y(x);
        while (k > 0) {
            if (k & 1) ret *= y;
            y *= y;
            k >>= 1;
        }
        return ret;
    }

    friend ostream &operator<<(ostream &os, const Arbitrary_Modint &p) {
        return os << p.x;
    }

    friend istream &operator>>(istream &is, Arbitrary_Modint &p) {
        int64_t y;
        is >> y;
        p = Arbitrary_Modint(y);
        return (is);
    }

    static int get_mod() {
        return MOD;
    }
};
int Arbitrary_Modint::MOD = 998244353;

using modint9 = Modint<998244353>;
using modint1 = Modint<1000000007>;
using modint  = Arbitrary_Modint;

// end other/Modint.hpp
// restart A.cpp
using mint = modint9;
// start math/Combination.hpp

template <typename T>
struct Combination {
    int N;
    vector<T> fact, invfact;
    Combination(int N) : N(N) {
        fact.resize(N + 1);
        invfact.resize(N + 1);
        fact[0] = 1;
        for (int i = 1; i <= N; i++) {
            fact[i] = fact[i - 1] * i;
        }
        invfact[N] = T(1) / fact[N];
        for (int i = N - 1; i >= 0; i--) {
            invfact[i] = invfact[i + 1] * (i + 1);
        }
    }

    void extend(int n) {
        int le = fact.size();
        fact.resize(n + 1);
        invfact.resize(n + 1);
        for (int i = le; i <= n; i++) {
            fact[i] = fact[i - 1] * i;
        }
        invfact[n] = T(1) / fact[n];
        for (int i = n - 1; i >= le; i--) {
            invfact[i] = invfact[i + 1] * (i + 1);
        }
    }

    T nCk(int n, int k) {
        if (k > n || k < 0) return T(0);
        if (n >= fact.size()) extend(n);
        return fact[n] * invfact[k] * invfact[n - k];
    }

    T nPk(int n, int k) {
        if (k > n || k < 0) return T(0);
        if (n >= fact.size()) extend(n);
        return fact[n] * invfact[n - k];
    }

    T nHk(int n, int k) {
        if (n == 0 && k == 0) return T(1);
        return nCk(n + k - 1, k);
    }

    T Catalan(int n) {
        return nCk(2 * n, n) - nCk(2 * n, n + 1);
    }

    // n 個の +1, m 個の -1, 累積和が常にk以下
    T Catalan(int n, int m, int k) {
        if (n > m + k || k < 0)
            return T(0);
        else
            return nCk(n + m, n) - nCk(n + m, m + k + 1);
    }
};
// end math/Combination.hpp
// restart A.cpp

void solve() {
    INT(n);
    Combination<mint> Comb(n + 10);
    auto edges = read_tree(n);
    vec(int, siz, n, 1);
    vec(int, leaf, n, 0);
    vec(int, tot, n + 1, 0);
    vec(int, cnt, n + 1, 0);

    int root = 0;
    fori(pos, n) {
        if (edges[pos].size() == 1) {
            root = pos;
            break;
        }
    }

    auto dfs = [&](auto &&self, int pos, int bpos) -> void {
        if (len(edges[pos]) == 1) leaf[pos] = 1;
        for (auto npos : edges[pos]) {
            if (npos == bpos) continue;
            self(self, npos, pos);
            siz[pos] += siz[npos];
            leaf[pos] += leaf[npos];
        }
    };

    dfs(dfs, root, -1);
    int le = leaf[root];

    fori(i, n) {
        tot[siz[i]] += leaf[i];
        tot[n - siz[i]] += le - leaf[i];

        cnt[siz[i]]++;
        cnt[n - siz[i]]++;
    }

    set<int> ng;
    fori(pos, n) {
        for (auto npos : edges[pos]) {
            if (siz[npos] + 1 == siz[pos]) {
                ng.insert(siz[npos]);
                ng.insert(n - siz[pos]);
            }
        }
    }

    mint ans = 0;
    fori(j, n + 1) {
        if (cnt[j] == 0) {
            continue;
        }

        int i = min(j, n - j);
        ans += Comb.nCk(n, i);
        if (0 < i && i < n && cnt[i] == 1 && (!ng.count(i) || cnt[i + 1] == 0 || 2 * i + 1 == n) && (cnt[i - 1] == 0) && cnt[i + 1] <= 1) {
            int a = tot[i];
            int b = le - tot[i];
            ans -= Comb.nCk(n - a - b, i - b);
        }
        int a = tot[i];
        int b = le - tot[i];
    }
    print(ans);
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    // cout << fixed << setprecision(12);
    int t;
    t = 1;
    // cin >> t;
    while (t--) solve();
    return 0;
}

// end A.cpp
yukicoder

結果

テストケース

ソースコード
