#1194980 (C++23) No.828 全方位神童数

• 問題: No.828 全方位神童数
• ユーザー: kwm_t
• 提出日時: 2026-10-24 15:15:26
• 言語: C++23 (gcc 15.3.0 + boost 1.92.0 + ACL)
• 結果: AC
• 実行時間: 331 ms / 2,000 ms
• コード長: 3,397 bytes
• コンパイル時間: 4,962 ms
• コンパイル使用メモリ: 387,352 KB
• 実行使用メモリ: 41,856 KB
• 最終ジャッジ日時: 2026-10-24 15:15:42
• 合計ジャッジ時間: 13,612 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge2_0

ファイルパターン	結果
sample	AC * 3
other	AC * 44

ソースコード

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using mint = modint1000000007;
const int mod = 1000000007;

#define rep(i, n) for (int i = 0; i < (n); ++i)
#define all(x) (x).begin(), (x).end()
#define endl "\n"

struct Treap {
    struct Node {
        int key, priority, size;
        int left = -1, right = -1;
    };

    vector<Node> t;
    int root = -1;

    Treap(int n) : t(n) {
        rep(i, n) {
            unsigned int x = i + 1;
            x += 0x9e3779b9;
            x = (x ^ (x >> 16)) * 0x85ebca6b;
            x = (x ^ (x >> 13)) * 0xc2b2ae35;
            x ^= x >> 16;
            t[i] = { i, (int)x, 1, -1, -1 };
        }
    }

    int size(int v) const {
        return v == -1 ? 0 : t[v].size;
    }

    void update(int v) {
        if (v == -1) return;
        t[v].size = 1 + size(t[v].left) + size(t[v].right);
    }

    // key 未満と key 以上に分割
    void split(int v, int key, int& a, int& b) {
        if (v == -1) {
            a = b = -1;
            return;
        }

        if (t[v].key < key) {
            a = v;
            split(t[v].right, key, t[v].right, b);
            update(v);
        }
        else {
            b = v;
            split(t[v].left, key, a, t[v].left);
            update(v);
        }
    }

    int merge(int a, int b) {
        if (a == -1) return b;
        if (b == -1) return a;

        if (t[a].priority > t[b].priority) {
            t[a].right = merge(t[a].right, b);
            update(a);
            return a;
        }
        else {
            t[b].left = merge(a, t[b].left);
            update(b);
            return b;
        }
    }

    // key 未満の要素を一時退避し、key 以上だけ残す
    int removeLess(int key) {
        int a, b;
        split(root, key, a, b);
        root = b;
        return a;
    }

    void restore(int saved) {
        root = merge(saved, root);
    }

    void insert(int v) {
        t[v].left = t[v].right = -1;
        t[v].size = 1;

        int a, b;
        split(root, v, a, b);
        root = merge(merge(a, v), b);
    }

    int count() const {
        return size(root);
    }
};

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

    int n;
    cin >> n;

    vector<int> r(n);
    rep(i, n) cin >> r[i];

    vector<vector<int>> to(n);
    rep(i, n - 1) {
        int a, b;
        cin >> a >> b;
        --a;
        --b;
        to[a].push_back(b);
        to[b].push_back(a);
    }

    Treap st(n);
    vector<long long> s(n);

    auto dfs = [&](auto self, int v, int p = -1) -> void {
        int saved = st.removeLess(v);

        st.insert(v);

        for (int nv : to[v]) {
            if (nv == p) continue;

            self(self, nv, v);
            st.removeLess(v);
        }

        s[v] += st.count();
        st.restore(saved);
    };

    dfs(dfs, 0);

    rep(i, n) reverse(all(to[i]));

    st.root = -1;

    auto dfs2 = [&](auto self, int v, int p = -1) -> void {
        int saved = st.removeLess(v);

        s[v] += st.count();

        for (int nv : to[v]) {
            if (nv == p) continue;

            self(self, nv, v);
            st.removeLess(v);
        }

        st.restore(saved);
        st.insert(v);
    };

    dfs2(dfs2, 0);

    mint ans = 1;
    rep(i, n) ans *= s[i] + r[i];

    cout << ans.val() << endl;
    return 0;
}