結果

問題 No.1002 Twotone
ユーザー yosupotyosupot
提出日時 2020-02-28 22:37:47
言語 C++17
(gcc 13.2.0 + boost 1.83.0)
結果
AC  
実行時間 963 ms / 5,000 ms
コード長 7,403 bytes
コンパイル時間 2,217 ms
コンパイル使用メモリ 152,004 KB
実行使用メモリ 61,340 KB
最終ジャッジ日時 2024-04-21 19:46:32
合計ジャッジ時間 14,930 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 2 ms
5,376 KB
testcase_02 AC 2 ms
5,376 KB
testcase_03 AC 309 ms
20,148 KB
testcase_04 AC 426 ms
25,132 KB
testcase_05 AC 410 ms
25,240 KB
testcase_06 AC 2 ms
5,376 KB
testcase_07 AC 207 ms
16,928 KB
testcase_08 AC 368 ms
25,944 KB
testcase_09 AC 367 ms
25,968 KB
testcase_10 AC 2 ms
5,376 KB
testcase_11 AC 471 ms
21,048 KB
testcase_12 AC 644 ms
25,984 KB
testcase_13 AC 657 ms
25,984 KB
testcase_14 AC 2 ms
5,376 KB
testcase_15 AC 291 ms
16,592 KB
testcase_16 AC 626 ms
26,156 KB
testcase_17 AC 627 ms
26,232 KB
testcase_18 AC 2 ms
5,376 KB
testcase_19 AC 490 ms
25,984 KB
testcase_20 AC 593 ms
33,280 KB
testcase_21 AC 615 ms
32,900 KB
testcase_22 AC 2 ms
5,376 KB
testcase_23 AC 461 ms
20,960 KB
testcase_24 AC 902 ms
32,596 KB
testcase_25 AC 896 ms
32,492 KB
testcase_26 AC 2 ms
5,376 KB
testcase_27 AC 53 ms
17,736 KB
testcase_28 AC 109 ms
24,640 KB
testcase_29 AC 86 ms
24,644 KB
testcase_30 AC 2 ms
5,376 KB
testcase_31 AC 78 ms
24,352 KB
testcase_32 AC 112 ms
24,616 KB
testcase_33 AC 82 ms
24,640 KB
testcase_34 AC 963 ms
61,340 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

//#pragma GCC optimize("Ofast")
//#pragma GCC target("avx")
//#undef LOCAL




#include <algorithm>

#include <array>

#include <bitset>

#include <cassert>

#include <complex>

#include <cstdio>

#include <cstring>

#include <iostream>

#include <map>

#include <numeric>

#include <queue>

#include <set>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include <vector>

using namespace std;

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;
constexpr ll TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }
template <class T> using V = vector<T>;
template <class T> using VV = V<V<T>>;

struct Scanner {
    FILE* fp = nullptr;
    char line[(1 << 15) + 1];
    size_t st = 0, ed = 0;
    void reread() {
        memmove(line, line + st, ed - st);
        ed -= st;
        st = 0;
        ed += fread(line + ed, 1, (1 << 15) - ed, fp);
        line[ed] = '\0';
    }
    bool succ() {
        while (true) {
            if (st == ed) {
                reread();
                if (st == ed) return false;
            }
            while (st != ed && isspace(line[st])) st++;
            if (st != ed) break;
        }
        if (ed - st <= 50) reread();
        return true;
    }
    template <class T, enable_if_t<is_same<T, string>::value, int> = 0>
    bool read_single(T& ref) {
        if (!succ()) return false;
        while (true) {
            size_t sz = 0;
            while (st + sz < ed && !isspace(line[st + sz])) sz++;
            ref.append(line + st, sz);
            st += sz;
            if (!sz || st != ed) break;
            reread();
        }
        return true;
    }
    template <class T, enable_if_t<is_integral<T>::value, int> = 0>
    bool read_single(T& ref) {
        if (!succ()) return false;
        bool neg = false;
        if (line[st] == '-') {
            neg = true;
            st++;
        }
        ref = T(0);
        while (isdigit(line[st])) {
            ref = 10 * ref + (line[st++] - '0');
        }
        if (neg) ref = -ref;
        return true;
    }
    template <class T> bool read_single(V<T>& ref) {
        for (auto& d : ref) {
            if (!read_single(d)) return false;
        }
        return true;
    }
    void read() {}
    template <class H, class... T> void read(H& h, T&... t) {
        bool f = read_single(h);
        assert(f);
        read(t...);
    }
    Scanner(FILE* _fp) : fp(_fp) {}
};

struct Printer {
  public:
    template <bool F = false> void write() {}
    template <bool F = false, class H, class... T>
    void write(const H& h, const T&... t) {
        if (F) write_single(' ');
        write_single(h);
        write<true>(t...);
    }
    template <class... T> void writeln(const T&... t) {
        write(t...);
        write_single('\n');
    }

    Printer(FILE* _fp) : fp(_fp) {}
    ~Printer() { flush(); }

  private:
    static constexpr size_t SIZE = 1 << 15;
    FILE* fp;
    char line[SIZE], small[50];
    size_t pos = 0;
    void flush() {
        fwrite(line, 1, pos, fp);
        pos = 0;
    }
    void write_single(const char& val) {
        if (pos == SIZE) flush();
        line[pos++] = val;
    }
    template <class T, enable_if_t<is_integral<T>::value, int> = 0>
    void write_single(T val) {
        if (pos > (1 << 15) - 50) flush();
        if (val == 0) {
            write_single('0');
            return;
        }
        if (val < 0) {
            write_single('-');
            val = -val; // todo min
        }
        size_t len = 0;
        while (val) {
            small[len++] = char('0' + (val % 10));
            val /= 10;
        }
        for (size_t i = 0; i < len; i++) {
            line[pos + i] = small[len - 1 - i];
        }
        pos += len;
    }
    void write_single(const string& s) {
        for (char c : s) write_single(c);
    }
    void write_single(const char* s) {
        size_t len = strlen(s);
        for (size_t i = 0; i < len; i++) write_single(s[i]);
    }
    template <class T> void write_single(const V<T>& val) {
        auto n = val.size();
        for (size_t i = 0; i < n; i++) {
            if (i) write_single(' ');
            write_single(val[i]);
        }
    }
};


struct Centroid {
    int r;
    VV<int> tr;
    V<int> par;
};

template <class E> struct CentroidExec : Centroid {
    int n;
    const VV<E>& g;
    V<bool> used;

    using P = pair<int, int>;
    V<P> info; //(child max, child)

    int dfs(int p, int b) {
        int sz = 1;
        info[p] = P(0, -1);
        for (E e : g[p]) {
            int d = e.to;
            if (d == b || used[d]) continue;
            int csz = dfs(d, p);
            sz += csz;
            info[p] = max(info[p], P(csz, d));
        }
        return sz;
    }
    int init(int p, int b) {
        int sz = dfs(p, -1);
        while (info[p].first > sz / 2) p = info[p].second;
        par[p] = b;
        used[p] = true;
        for (E e : g[p]) {
            int d = e.to;
            if (used[d]) continue;
            tr[p].push_back(init(d, p));
        }
        return p;
    }
    CentroidExec(const VV<E>& _g) : n(int(_g.size())), g(_g), used(n), info(n) {
        tr = VV<int>(n);
        par = V<int>(n);
        r = init(0, -1);
    }
};

template <class E> Centroid get_centroid(const VV<E>& g) {
    return CentroidExec<E>(g);
}

Scanner sc = Scanner(stdin);
Printer pr = Printer(stdout);

int main() {
    int n, k;
    sc.read(n, k);
    struct E {
        int to, col;
    };
    VV<E> g(n);
    for (int i = 0; i < n - 1; i++) {
        int a, b, c;
        sc.read(a, b, c); a--; b--;
        g[a].push_back({b, c});
        g[b].push_back({a, c});
    }
    auto cent = get_centroid(g);

    V<bool> cvis(n);
    using P = pair<int, int>;
    ll ans = 0;
    auto dfs = [&](auto self, int p) -> void {
        int mp1_sum = 0;
        map<int, int> mp1, mp2_sum;
        map<P, int> mp2;
        V<P> buf;
        auto inc = [&](int c0, int c1) {
            if (c0 < c1) swap(c0, c1);
            if (c1 != -1) {
                ans++;
                ans += mp2[{c0, c1}] + mp1[c0] + mp1[c1];
            } else {
                ans += mp1_sum - mp1[c0] + mp2_sum[c0];
            }
            buf.push_back({c0, c1});
        };
        auto ref = [&]() {
            for (auto p: buf) {
                if (p.second != -1) {
                    mp2[p]++;
                    mp2_sum[p.first]++;
                    mp2_sum[p.second]++;
                } else {
                    mp1[p.first]++;
                    mp1_sum++;
                }
            }
            buf.clear();
        };
        auto dfs2 = [&](auto self, int p, int b, int c0, int c1) -> void {
            inc(c0, c1);
            for (auto e: g[p]) {
                int d = e.to;
                int c = e.col;
                if (d == b || cvis[d]) continue;
                if (c == c0 || c == c1) {
                    self(self, d, p, c0, c1);
                } else if (c1 == -1) {
                    self(self, d, p, c0, c);
                }
            }
        };
        for (auto e: g[p]) {
            int d = e.to;
            if (cvis[d]) continue;
            dfs2(dfs2, d, p, e.col, -1);
            ref();
        }

        cvis[p] = true;
        for (auto d: cent.tr[p]) {
            self(self, d);
        }
    };
    dfs(dfs, cent.r);
    pr.writeln(ans);
    return 0;
}
0