#436358 (C++17) No.1002 Twotone

提出ソース
結果

問題	No.1002 Twotone
ユーザー	Pachicobue
提出日時	2020-02-28 23:22:59
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	WA
実行時間	-
コード長	12,691 bytes
コンパイル時間	3,451 ms
コンパイル使用メモリ	248,032 KB
実行使用メモリ	82,076 KB
最終ジャッジ日時	2024-10-13 19:13:28
合計ジャッジ時間	13,715 ms
ジャッジサーバーID （参考情報）	judge2 / judge5
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	WA	-
testcase_01	WA	-
testcase_02	WA	-
testcase_03	WA	-
testcase_04	WA	-
testcase_05	WA	-
testcase_06	WA	-
testcase_07	TLE	-
testcase_08	--	-
testcase_09	--	-
testcase_10	--	-
testcase_11	--	-
testcase_12	--	-
testcase_13	--	-
testcase_14	--	-
testcase_15	--	-
testcase_16	--	-
testcase_17	--	-
testcase_18	--	-
testcase_19	--	-
testcase_20	--	-
testcase_21	--	-
testcase_22	--	-
testcase_23	--	-
testcase_24	--	-
testcase_25	--	-
testcase_26	--	-
testcase_27	--	-
testcase_28	--	-
testcase_29	--	-
testcase_30	--	-
testcase_31	--	-
testcase_32	--	-
testcase_33	--	-
testcase_34	--	-
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include <bits/stdc++.h>
// created [2020/02/28] 21:45:15
#pragma GCC diagnostic ignored "-Wsign-compare"
#pragma GCC diagnostic ignored "-Wsign-conversion"

using i32   = int32_t;
using i64   = int64_t;
using u32   = uint32_t;
using u64   = uint64_t;
using uint  = unsigned int;
using usize = std::size_t;
using ll    = long long;
using ull   = unsigned long long;
using ld    = long double;
template<typename T, usize n>
using arr = T (&)[n];
template<typename T, usize n>
using c_arr = const T (&)[n];
template<typename T>
using max_heap = std::priority_queue<T>;
template<typename T>
using min_heap = std::priority_queue<T, std::vector<T>, std::greater<T>>;
template<typename T> constexpr T popcount(const T u) { return u ? static_cast<T>(__builtin_popcountll(static_cast<u64>(u))) : static_cast<T>(0); }
template<typename T> constexpr T log2p1(const T u) { return u ? static_cast<T>(64 - __builtin_clzll(static_cast<u64>(u))) : static_cast<T>(0); }
template<typename T> constexpr T msbp1(const T u) { return log2p1(u); }
template<typename T> constexpr T lsbp1(const T u) { return __builtin_ffsll(u); }
template<typename T> constexpr T clog(const T u) { return u ? log2p1(u - 1) : static_cast<T>(u); }
template<typename T> constexpr bool ispow2(const T u) { return u and (static_cast<u64>(u) & static_cast<u64>(u - 1)) == 0; }
template<typename T> constexpr T ceil2(const T u) { return static_cast<T>(1) << clog(u); }
template<typename T> constexpr T floor2(const T u) { return u == 0 ? static_cast<T>(0) : static_cast<T>(1) << (log2p1(u) - 1); }
template<typename T> constexpr bool btest(const T mask, const usize ind) { return static_cast<bool>((static_cast<u64>(mask) >> ind) & static_cast<u64>(1)); }
template<typename T> void bset(T& mask, const usize ind) { mask |= (static_cast<T>(1) << ind); }
template<typename T> void breset(T& mask, const usize ind) { mask &= ~(static_cast<T>(1) << ind); }
template<typename T> void bflip(T& mask, const usize ind) { mask ^= (static_cast<T>(1) << ind); }
template<typename T> void bset(T& mask, const usize ind, const bool b) { (b ? bset(mask, ind) : breset(mask, ind)); }
template<typename T> constexpr T bcut(const T mask, const usize ind) { return ind == 0 ? static_cast<T>(0) : static_cast<T>((static_cast<u64>(mask) << (64 - ind)) >> (64 - ind)); }
template<typename T> bool chmin(T& a, const T& b) { return (a > b ? a = b, true : false); }
template<typename T> bool chmax(T& a, const T& b) { return (a < b ? a = b, true : false); }
constexpr unsigned int mod                  = 1000000007;
template<typename T> constexpr T inf_v      = std::numeric_limits<T>::max() / 4;
template<typename Real> constexpr Real pi_v = Real{3.141592653589793238462643383279502884};
auto mfp = [](auto&& f) { return [=](auto&&... args) { return f(f, std::forward<decltype(args)>(args)...); }; };

template<typename T>
T in()
{
    T v;
    return std::cin >> v, v;
}
template<typename T, typename Uint, usize n, usize i>
T in_v(typename std::enable_if<(i == n), c_arr<Uint, n>>::type) { return in<T>(); }
template<typename T, typename Uint, usize n, usize i>
auto in_v(typename std::enable_if<(i < n), c_arr<Uint, n>>::type& szs)
{
    const usize s = (usize)szs[i];
    std::vector<decltype(in_v<T, Uint, n, i + 1>(szs))> ans(s);
    for (usize j = 0; j < s; j++) { ans[j] = in_v<T, Uint, n, i + 1>(szs); }
    return ans;
}
template<typename T, typename Uint, usize n>
auto in_v(c_arr<Uint, n> szs) { return in_v<T, Uint, n, 0>(szs); }
template<typename... Types>
auto in_t() { return std::tuple<std::decay_t<Types>...>{in<Types>()...}; }
struct io_init
{
    io_init()
    {
        std::cin.tie(nullptr), std::ios::sync_with_stdio(false);
        std::cout << std::fixed << std::setprecision(20);
    }
    void clear()
    {
        std::cin.tie(), std::ios::sync_with_stdio(true);
    }
} io_setting;

int out() { return 0; }
template<typename T>
int out(const T& v) { return std::cout << v, 0; }
template<typename T>
int out(const std::vector<T>& v)
{
    for (usize i = 0; i < v.size(); i++) {
        if (i > 0) { std::cout << ' '; }
        out(v[i]);
    }
    return 0;
}
template<typename T1, typename T2>
int out(const std::pair<T1, T2>& v) { return out(v.first), std::cout << ' ', out(v.second), 0; }
template<typename T, typename... Args>
int out(const T& v, const Args... args) { return out(v), std::cout << ' ', out(args...), 0; }
template<typename... Args>
int outln(const Args... args) { return out(args...), std::cout << '\n', 0; }
template<typename... Args>
int outel(const Args... args) { return out(args...), std::cout << std::endl, 0; }
#    define SHOW(...) static_cast<void>(0)
constexpr ull TEN(const usize n) { return n == 0 ? 1ULL : TEN(n - 1) * 10ULL; }

template<typename T, typename Uint, usize n, usize i>
auto make_v(typename std::enable_if<(i == n), c_arr<Uint, n>>::type, const T& v = T{}) { return v; }
template<typename T, typename Uint, usize n, usize i>
auto make_v(typename std::enable_if<(i < n), c_arr<Uint, n>>::type szs, const T& v = T{})
{
    const usize s = (usize)szs[i];
    return std::vector<decltype(make_v<T, Uint, n, i + 1>(szs, v))>(s, make_v<T, Uint, n, i + 1>(szs, v));
}
template<typename T, typename Uint, usize n>
auto make_v(c_arr<Uint, n> szs, const T& t = T{}) { return make_v<T, Uint, n, 0>(szs, t); }

template<typename Cost = usize>
struct edge
{
    using cost_type = Cost;
    usize u, v;
    Cost c;
    edge(const usize u, const usize v) : u{u}, v{v}, c{1} {}
    edge(const usize u, const usize v, const Cost& c) : u{u}, v{v}, c{c} {}
    operator usize() const { return v; }
    usize from() const { return u; }
    usize to() const { return v; }
    Cost cost() const { return c; }
    friend std::ostream& operator<<(std::ostream& os, const edge& e) { return os << e.u << "->" << e.v << ":" << e.c; }
};
template<typename Edge>
class base_graph
{
public:
    base_graph(const usize n) : v{n}, es(n), res(n) {}
    void add_edge(const usize u, const usize v, const bool bi = false)
    {
        es[u].emplace_back(u, v), res[v].emplace_back(v, u);
        if (bi) { es[v].emplace_back(v, u), res[u].emplace_back(u, v); }
    }
    template<typename Cost>
    void add_edge(const usize u, const usize v, const Cost& c, const bool bi = false)
    {
        es[u].emplace_back(u, v, c), res[v].emplace_back(v, u, c);
        if (bi) { es[v].emplace_back(v, u, c), res[u].emplace_back(u, v, c); }
    }
    std::vector<Edge>& operator[](const usize u) { return es[u]; }
    const std::vector<Edge>& operator[](const usize u) const { return es[u]; }
    std::vector<Edge>& from(const usize u) { return es[u]; }
    const std::vector<Edge>& from(const usize u) const { return es[u]; }
    std::vector<Edge>& to(const usize v) { return res[v]; }
    const std::vector<Edge>& to(const usize v) const { return res[v]; }
    usize size() const { return v; }
    friend std::ostream& operator<<(std::ostream& os, const base_graph& g)
    {
        for (usize i = 0; i < g.v; i++) {
            for (const auto& e : g.es[i]) { os << e << '\n'; }
        }
        return os;
    }

private:
    usize v;
    std::vector<std::vector<Edge>> es, res;
};
template<typename Edge>
using base_tree = base_graph<Edge>;
using graph     = base_graph<edge<>>;
using tree      = base_graph<edge<>>;
template<typename Cost>
using cost_graph = base_graph<edge<Cost>>;
template<typename Cost>
using cost_tree = base_graph<edge<Cost>>;
template<typename Edge = edge<>>
class centroid
{
public:
    centroid(const base_tree<Edge>& tree) : cs(tree.size())
    {
        const std::size_t sz = tree.size();
        std::vector<std::size_t> sub(sz, 1);
        std::vector<bool> used(sz, false);
        auto size = [&, sz](auto&& self, const std::size_t s, const std::size_t p) -> std::size_t {
            sub[s] = 1;
            for (const auto& e : tree[s]) {
                const std::size_t to = e.to();
                if (to == p or used[to]) { continue; }
                sub[s] += self(self, to, s);
            }
            return sub[s];
        };
        auto search = [&, sz](auto&& self, const std::size_t s, const std::size_t p, const std::size_t tot) -> std::size_t {
            for (const auto& e : tree[s]) {
                const std::size_t to = e.to();
                if (p == to or used[to]) { continue; }
                if (sub[to] * 2 > tot) { return self(self, to, s, tot); }
            }
            return s;
        };
        auto build = [&, sz](auto&& self, const std::size_t s, const std::size_t pc) -> std::size_t {
            const std::size_t tot = size(size, s, sz), c = search(search, s, sz, tot);
            used[c] = true;
            if (pc != sz) { cs.add_edge(pc, c); }
            for (const auto& e : tree[c]) {
                const std::size_t to = e.to();
                if (not used[to]) { self(self, to, c); }
            }
            return c;
        };
        build(build, 0, sz);
    }
    const tree& centros() const { return cs; }

private:
    tree cs;
};
int main()
{
    const auto N = in<int>();
    const auto K = in<int>();
    cost_graph<int> g(N);
    for (int i = 0; i < N - 1; i++) {
        const auto u = in<int>() - 1, v = in<int>() - 1, c = in<int>() - 1;
        g.add_edge(u, v, c, true);
    }
    const auto cs = centroid(g).centros();
    int c         = 0;
    for (; c < N; c++) {
        if (cs.to(c).empty()) { break; }
    }
    SHOW(cs);
    std::vector<bool> used(N, false);
    ll ans = 0;
    mfp([&](auto&& self, const int c) -> void {
        used[c] = true;
        std::map<std::set<int>, int> total;              // [A,B,C]->個数
        std::vector<std::map<std::set<int>, int>> subs;  // 部分木ごとの [A,B,C]->個数
        total[std::set<int>{}]++;
        for (const auto& e : g[c]) {
            subs.push_back(std::map<std::set<int>, int>{});
            const int ncolor = e.cost();
            const int ns     = e.to();
            if (used[ns]) { continue; }
            mfp([&](auto&& self, const int s, const int p, const std::set<int>& cols) -> void {
                total[cols]++;
                subs.back()[cols]++;
                for (const auto& e : g[s]) {
                    const int to    = e.to();
                    const int color = e.cost();
                    if (to == p) { continue; }
                    std::set<int> ncols = cols;
                    ncols.insert(color);
                    if (ncols.size() > 2) { continue; }
                    self(self, to, s, ncols);
                }
            })(ns, c, std::set<int>{ncolor});
        }
        SHOW(c);
        SHOW(total);
        SHOW(subs);
        ll sum = 0;
        std::map<int, int> total_monos;                          // [A?]->個数
        std::vector<std::map<int, int>> sub_monos(subs.size());  // 部分木ごとの [A?]->個数
        for (const auto& p : total) {
            const auto& st = p.first;
            for (const int c : st) { total_monos[c] += p.second; }
        }
        for (const int nc : cs[c]) { self(self, nc); }
        for (int i = 0; i < subs.size(); i++) {
            for (const auto& p : subs[i]) {
                const auto& st = p.first;
                for (const int c : st) { sub_monos[i][c] += p.second; }
            }
        }
        auto count = [&](const std::map<std::set<int>, int>& mp, const std::map<int, int>& mono) {
            ll zero = 0;
            ll one  = 0;
            ll two  = 0;
            for (const auto& p : mp) {
                const auto& st = p.first;
                const int cnt  = p.second;
                if (st.size() == 0) {
                    zero += cnt;
                } else if (st.size() == 1) {
                    one += cnt;
                } else if (st.size() == 2) {
                    two += cnt;
                } else {
                    SHOW("HOGEHOGE");
                }
            }
            ll ans = 0;
            for (const auto& p : mp) {
                const auto& st = p.first;
                const int cnt  = p.second;
                if (st.size() == 0) {
                    ans += two + one + zero;
                } else if (st.size() == 1) {
                    ans += (mono.count(*st.begin()) ? mono.at(*st.begin()) : 0LL) - cnt;
                } else if (st.size() == 2) {
                    ans += cnt;
                } else {
                    SHOW("HOGEHOGE");
                }
                for (const int c : st) {
                }
            }
            return ans;
        };
        ans += count(total, total_monos);
        for (int i = 0; i < subs.size(); i++) { ans -= count(subs[i], sub_monos[i]); }
    })(c);
    outln(ans);
    return 0;
}
yukicoder

結果

テストケース

ソースコード
