Submission #1181949

#	Time	Username	Problem	Language	Result	Execution time	Memory
1181949		dpsaveslives	Wells (CEOI21_wells)	C++20	100 / 100	3516 ms	303908 KiB

wells

#include <bits/stdc++.h>
#define For(i,l,r) for(register int i=(l);i<=(r);++i)
#define For_down(i,r,l) for(register int i=(r);i>=(l);--i)
using namespace std;
constexpr int n_MAX = 1500000 + 5, mod = 1000000007;
int norm(const int x) {
    return ((x >= mod) ? (x - mod) : x);
}
void add(int &x, const int y) {
    (((x += y) >= mod) && (x -= mod));
}
int mul(const int x, const int y) {
    return (((long long)x * y) % mod);
}
int n, k, p, q, diam, Fa[n_MAX], dep[n_MAX], dis[n_MAX], other[n_MAX], project[n_MAX], h[n_MAX], side[n_MAX],
    equ_class[n_MAX];
vector<int> A[n_MAX], ch[n_MAX], path;
void DFS(const int u, int dis[]) {
    for (const int v : A[u])
        if (v != Fa[u])
            Fa[v] = u, dis[v] = (dis[u] + 1), DFS(v, dis);
}
int max_element() {
    return (max_element(dis + 1, dis + n + 1) - dis);
}
void DFS2(const int u) {
    for (const int v : A[u])
        if (!project[v])
            Fa[v] = u, dep[v] = (dep[u] + 1), project[v] = project[u], DFS2(v), h[u] = max(h[u], h[v] + 1);
}
int tag[n_MAX];
bool ok[2][n_MAX];
int F[n_MAX], mul1[n_MAX], mul2[n_MAX], ans[2][n_MAX];
struct range {
    int l, r;
    bool c;
};
range norm(int l, int r) {
    if (l > r)
        return {0, k - 1, true};

    if ((r - l) >= (k - 1))
        return {0, k - 1};

    l %= k, r %= k;

    if (l <= r)
        return {l, r};
    else
        return {r + 1, l - 1, true};
}
range complement(const range a) {
    return {a.l, a.r, !a.c};
}
void ban(const int l, const int r) {
    if (l <= r)
        ++tag[l], --tag[r + 1];
}
void ban_intersect(range a, range b) {
    if (!a.c && !b.c)
        ban(max(a.l, b.l), min(a.r, b.r));
    else if (a.c && b.c)
        ban(0, min(a.l, b.l) - 1), ban(b.r + 1, a.l - 1), ban(a.r + 1, b.l - 1), ban(max(a.r, b.r) + 1, k - 1);
    else {
        if (a.c)
            swap(a, b);

        ban(a.l, min(a.r, b.l - 1)), ban(max(a.l, b.r + 1), a.r);
    }
}
void ban(const int l1, const int r1, const int l2, const int r2) {
    const range a = norm(l1, r1), b = norm(l2, r2);
    ban_intersect(a, b), ban_intersect(complement(a), complement(b));
}
void ban(const int l, const int r, const int p) {
    ban_intersect(complement(norm(l, r)), {p, p, true});
}
void apply(const int l1, int r1, const int l2, const int r2, const int x) {
    if (l2 == (r1 + 1)) {
        if (r2 < (k - 1))
            mul1[r2 + 1] = mul(mul1[r2 + 1], x);
    } else {
        if ((r2 - l2) >= (k - 1))
            return;

        r1 = max(r1, r2 - k), mul2[r1 + 1] = mul(mul2[r1 + 1], x);

        if (r2 < (k - 1))
            mul1[r2 + 1] = mul(mul1[r2 + 1], x);
    }
}
void DP(const int u) {
    if (!ch[u].empty()) {
        F[u] = 1;

        for (const int v : ch[u])
            DP(v), F[u] = mul(F[u], F[v]);
    }

    if (dis[u] <= (k - 1))
        add(F[u], 1);
}
void DFS3(const int u) {
    const int dis_proj = dis[project[u]], l = (max(dis_proj, (k - 1) >> 1) + 1), r = dis[u], len2 = (r - l + 1),
              len1 = (k - len2);

    if (dep[u] <= (k - 1)) {
        if (dis[u] >= (k - 1))
            ban(dis_proj - len1 + 1, dis_proj, l, r);

        if (other[u] >= (k - 1))
            ban(dis_proj, dis_proj + len1 - 1, l, r);
    }

    for (const int v : ch[u])
        if (side[v] == 1)
            DP(v), apply(0, dis_proj, max(dis_proj, (k - 1) >> 1) + 1, dis[u], F[v]);
        else
            DFS3(v);
}
void DFS4(const int u) {
    int f = 0, s = 0;

    for (const int v : ch[u]) {
        DFS4(v);

        if ((h[v] + 1) > f)
            s = f, f = (h[v] + 1);
        else if ((h[v] + 1) > s)
            s = (h[v] + 1);
    }

    if ((f + s) >= (k - 1)) {
        const int l = min(s, (k - 1) >> 1), r = (k - l - 1);
        ban(max(dis[u] + l, (k - 1) >> 1) + 1, dis[u] + r, equ_class[u]);
    }
}
int main() {
    ios::sync_with_stdio(false), cin.tie(nullptr);
    cin >> n >> k;
    For(i, 1, n - 1) {
        int u, v;
        cin >> u >> v, A[u].push_back(v), A[v].push_back(u);
    }
    DFS(1, dis), p = max_element(), Fa[p] = dis[p] = 0, DFS(p, dis), q = max_element(), Fa[q] = 0, DFS(q, other),
                                            diam = dis[q];

    if (diam < (k - 1)) {
        int ans = 1;
        For(i, 1, n) ans = norm(ans << 1);
        cout << "YES\n" << ans << '\n';
        return 0;
    }

    int u = p;

    while (u)
        side[u] = 2, path.push_back(u), project[u] = u, u = Fa[u];

    For(i, 0, diam) {
        const int u = path[i];
        Fa[u] = dep[u] = 0, DFS2(u);
    }
    int cnt = 0;

    For(i, 1, n) if (!side[i]) {
        side[i] = (((dis[i] + h[i]) >= (k - 1)) + ((other[i] + h[i]) >= (k - 1)));

        if (!side[i])
            ++cnt;
    }

    For(i, 1, n) if (Fa[i] && side[i])
        ch[Fa[i]].push_back(i);

    For(I, 0, 1) {
        memset(tag, 0, k << 2), fill(mul1, mul1 + k, 1), fill(mul2, mul2 + k, 1);

        For(i, 1, n) if (side[i] == 2) {
            if (dep[i] >= ((k - 1) >> 1))
                equ_class[i] = (dis[i] % k);
            else
                equ_class[i] = (-1);
        }

        For_down(i, diam, (diam + I + 1) >> 1) {
            const int u = path[i];
            DFS3(u);

            for (const int v : ch[u])
                DFS4(v);
        }
        For(i, 0, k - 1) {
            if (!tag[i])
                ok[I][i] = true, ans[I][i] = 1;

            tag[i + 1] += tag[i];
        }
        For(i, 0, k - 1) ans[I][i] = mul(ans[I][i], mul1[i]), mul1[i + 1] = mul(mul1[i + 1], mul1[i]);
        For(i, 0, (k - 1) >> 1) ans[I][i] = mul(ans[I][i], mul2[i]), mul2[i + 1] = mul(mul2[i + 1], mul2[i]);
        swap(p, q);
        For(i, 1, n) swap(dis[i], other[i]);
        reverse(path.begin(), path.end());
    }
    bool Flag = false;
    int sum = 0;
    For(i, 0, k - 1) {
        const int j = ((diam - i) % k);
        Flag |= (ok[0][i] && ok[1][j]), add(sum, mul(ans[0][i], ans[1][j]));
    }
    For(i, 1, cnt) sum = norm(sum << 1);

    if (Flag)
        cout << "YES\n" << sum << '\n';
    else
        cout << "NO\n0\n";

    return 0;
}

Compilation message (stderr)

wells.cpp: In function 'int main()':
wells.cpp:141:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  141 |     For(i, 1, n - 1) {
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:150:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  150 |         For(i, 1, n) ans = norm(ans << 1);
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:160:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  160 |     For(i, 0, diam) {
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:166:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  166 |     For(i, 1, n) if (!side[i]) {
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:173:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  173 |     For(i, 1, n) if (Fa[i] && side[i])
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:176:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  176 |     For(I, 0, 1) {
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:179:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  179 |         For(i, 1, n) if (side[i] == 2) {
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:186:18: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  186 |         For_down(i, diam, (diam + I + 1) >> 1) {
      |                  ^
wells.cpp:3:42: note: in definition of macro 'For_down'
    3 | #define For_down(i,r,l) for(register int i=(r);i>=(l);--i)
      |                                          ^
wells.cpp:193:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  193 |         For(i, 0, k - 1) {
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:199:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  199 |         For(i, 0, k - 1) ans[I][i] = mul(ans[I][i], mul1[i]), mul1[i + 1] = mul(mul1[i + 1], mul1[i]);
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:200:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  200 |         For(i, 0, (k - 1) >> 1) ans[I][i] = mul(ans[I][i], mul2[i]), mul2[i + 1] = mul(mul2[i + 1], mul2[i]);
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:202:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  202 |         For(i, 1, n) swap(dis[i], other[i]);
      |             ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:207:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  207 |     For(i, 0, k - 1) {
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^
wells.cpp:211:9: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  211 |     For(i, 1, cnt) sum = norm(sum << 1);
      |         ^
wells.cpp:2:37: note: in definition of macro 'For'
    2 | #define For(i,l,r) for(register int i=(l);i<=(r);++i)
      |                                     ^

• Subtask 1: 30 / 30
• Subtask 2: 20 / 20
• Subtask 3: 20 / 20
• Subtask 4: 30 / 30