Submission #678694

#	Time	Username	Problem	Language	Result	Execution time	Memory
678694		stanislavpolyn	Wells (CEOI21_wells)	C++17	0 / 100	26 ms	23840 KiB

wells

#include <bits/stdc++.h>

#define fr(i, a, b) for (int i = (a); i <= (b); ++i)
#define rf(i, a, b) for (int i = (a); i >= (b); --i)
#define fe(x, y) for (auto& x : y)

#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define mt make_tuple

#define all(x) (x).begin(), (x).end()
#define pw(x) (1LL << (x))
#define sz(x) (int)(x).size()

using namespace std;

#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <typename T>
using oset = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
#define fbo find_by_order
#define ook order_of_key

template <typename T>
bool umn(T& a, T b) {
    return a > b ? a = b, 1 : 0;
}
template <typename T>
bool umx(T& a, T b) {
    return a < b ? a = b, 1 : 0;
}

using ll = long long;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
template <typename T>
using ve = vector<T>;

const int N = 1e6 + 5;
const int MOD = 1e9 + 7;

int add(int a, int b) {
    a += b;
    if (a >= MOD) {
        return a - MOD;
    }
    return a;
}

int mult(int a, int b) {
    return 1LL * a * b % MOD;
}

int n, k;
ve<int> g[N];
int dep[N];
int par[N];
int curRem;
bool have[N];
bool bad;
int mxDep[N];

void dfs(int v = 1, int p = 0) {
    par[v] = p;
    if (dep[v] % k == 0) {
        have[v] = 1;
    } else {
        have[v] = 0;
    }
    mxDep[v] = dep[v];
    int cnt = 0;

    pii mx = {-1e9, -1e9};
    fe (to, g[v]) {
        if (to == p) continue;

        dep[to] = dep[v] + 1;
        dfs(to, v);
        have[v] |= have[to];
        cnt += have[to];

        umx(mxDep[v], mxDep[to]);

        if (mx.fi < mxDep[to]) {
            mx.se = mx.fi;
            mx.fi = mxDep[to];
        } else {
            umx(mx.se, mxDep[to]);
        }
    }
//    cout << "GO " << v << " " << have[v] << "\n";
//    cout << mx.fi - dep[v] << " " << mx.se - dep[v] << "\n";
    if (!have[v] && (mx.fi - dep[v]) + (mx.se - dep[v]) + 1 >= k) {
//        cout << "CHECK " << v << "\n";
        bad = 1;
    }

    if (cnt > 1 && (2 * dep[v]) % k) {
        bad = 1;
    }
}

bool BAD = 0;
void dfs2(int v, int p, int dep, int cnt) {
    if (dep == k) {
//        assert(cnt == 1);
        if (cnt != 1) BAD = 1;
        return;
    }
    fe (to, g[v]) {
        if (to == p) continue;
        dfs2(to, v, dep + 1, cnt + (::dep[to] % k == 0));
    }
}

int main() {
#ifdef LOCAL
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
#else
    ios::sync_with_stdio(0);
    cin.tie(0);
#endif

    cin >> n >> k;
    fr (i, 1, n - 1) {
        int a, b;
        cin >> a >> b;

        g[a].pb(b);
        g[b].pb(a);
    }

    fr (root, 1, n) {
//        cout << "ROOT: " << root << "\n";
        bad = 0;
        dep[root] = 0;
        dfs(root, 0);

//        if (!bad) {
//            cout << "root: " << root << "\n";
            BAD = 0;
            fr (i, 1, n) {
                dfs2(i, 0, 1, dep[i] % k == 0);
            }
            if (!BAD) {
                cout << "YES\n";
                cout << "1\n";
                exit(0);
            }
//        }
    }

    cout << "NO\n";
//    cout << "1\n";

    return 0;
}

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