Submission #765092

#	Time	Username	Problem	Language	Result	Execution time	Memory
765092		drdilyor	Šarenlist (COCI22_sarenlist)	C++17	20 / 110	11 ms	312 KiB

Main

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
constexpr int inf = 1e9;
constexpr ll infl = 1e18;
constexpr ll mod = 1e9 + 7;

ll power(ll a, ll b) {
    ll res = 1;
    while (b-->0) res = res * a % mod;
    return res;
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    int n, m, k;
    cin >> n >> m >> k;
    vector<vector<pair<int,int>>> adj(n);
    for( int i = 0; i < n-1; i++) {
        int a, b;
        cin >> a >> b;
        a--;b--;
        adj[a].emplace_back(b, i);
        adj[b].emplace_back(a, i);
    }
    vector<pair<int,int>> p(m);
    for (auto& [i, j] : p) cin >> i >> j, i--,j--;

    auto path = [&](auto& path, int i, int j, int p=-1)->optional<vector<int>>{
        if (i == j) return vector<int>{};
        for (auto [e, ei] : adj[i]) {
            if (e == p) continue;
            optional<vector<int>> res = path(path, e, j, i);
            if (res.has_value()) {
                res.value().push_back(ei);
                return res;
            }
        }
        return nullopt;
    };


    auto sub4 = [&]() {
        int res = 0;
        for (int mask = 0; mask < (1<<(n-1)); mask++) {
            bool ok = 1;
            for (auto[i, j] : p) {
                vector edges = path(path, i, j).value();
                int black = 0, white = 0;
                for (int e : edges) {
                    if (mask & (1<<e)) black++;
                    else white++;
                }
                if (!black || !white) ok = 0;
                if (!ok) break;
            }
            res += ok;
        }
        cout << res << endl;
    };

    auto sub1 = [&]() {
        int len = path(path, p[0].first, p[0].second)->size();
        cout << (power(k, n-1) - k * power(k, n-1-len) % mod + mod) % mod << endl;
    };

    auto sub2 = [&]() {
        auto p1 = path(path, p[0].first, p[0].second).value();
        auto p2 = path(path, p[1].first, p[1].second).value();

        int len1 = p1.size();
        int len2 = p2.size();
        int over = 0;
        for (int i : p1) over += find(p2.begin(), p2.end(), i) != p2.end();

        ll res = power(k, n-1);
        res -= k * power(k, n-1 - len1) % mod;
        res -= k * power(k, n-1 - len2) % mod;
        res += k * power(k, n-1 - (len1 + len2 - over)) % mod;
        cout << (res % mod + mod) % mod << endl;
    };

    auto sub3 = [&]() {
        vector<int> len(m);
        for (int i = 0; i < m; i++)
            len[i] = path(path, p[i].first, p[i].second)->size();

        ll res = power(k, n - 1);
        for (int i = 0; i < m; i++) {
            res -= k * power(k, n-1 - len[i]) % mod;
            assert(len[i] > 1);
        }
        cout << (res + mod) % mod << '\n';
    };

    if (m == 1) sub1();
    // else if (m == 2) sub2();
    else if (n <= 15 && k == 2) sub4();
    else {
        sub3();
        // assert(false);
    }
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:67:10: warning: variable 'sub2' set but not used [-Wunused-but-set-variable]
   67 |     auto sub2 = [&]() {
      |          ^~~~

• Subtask 1: 10 / 10
• Subtask 2: 0 / 15
• Subtask 3: 0 / 10
• Subtask 4: 10 / 10
• Subtask 5: 0 / 65