5. Iba aby hladný nebol vzorové riešenie

5 4
1 3 0
2 3 0
3 4 1
4 5 0

528

3 3
1 2 0
3 2 0

0

Na vstupe sme mali strom. Našou úlohou bolo spočítať, koľko existuje postupností vrcholov dĺžky \(k\) takých, že počas prechádzania postupnosti prejdeme po aspoň jednej označenej hrane. Takéto postupnosti budeme nazývať dobrými.

Hrubá sila

Riešenie hrubou silou je priamočiare, nie však jednoduché. Nezľaknite sa preto, ak ho budete čítať. Vzorové riešenie sa ukáže ako omnoho jednoduchšie.

Môžeme si postupne vygenerovať všetky možné postupnosti dĺžky \(k\) a pre každú z nich overiť, či je dobrá, alebo nie.

Na implementáciu takéhoto riešenia vieme použiť napríklad rekurzívny backtracking. Keď máme postupnosť vygenerovanú, musíme overiť, či je dobrá. To vieme napríklad tak, že pre každý vrchol v postupnosti pustíme BFS, či DFS. V tomto prehľadávaní overíme, či na najkratšej ceste z aktuálneho vrchola postupnosti do nasledujúceho prejdeme po označenej hrane.

Inou, o niečo krajšou možnosťou by bolo predpočítať si pre každú dvojicu vrcholov, či sa na najkratšej ceste medzi nimi nachádza označená hrana. Potom by sme vedeli už počas rekurzie zisťovať, či sme cez nejakú označenú hranu prešli, alebo nie. Vyhli by sme sa tak overovaniu na konci rekurzie. Toto predpočítanie vieme spraviť drobnou úpravou algoritmu Floyda a Warshalla.

Možných postupností \(n\) vrcholov dĺžky \(k\) je \(n^k\). Každú z nich samostane vygenerujeme a overíme, či je dobrá. V závislosti od implementácie nám to bude trvať niečo medzi \(O(kn^{k + 1})\) a \(O(n^{k + 3})\). Pamäť bude tiež v závislosti od implementácie \(O(n)\) až \(O(n^2)\). Závisí to od toho, či si pamätáme iba strom, alebo si niečo predpočítame pre každú dvojicu vrcholov.

Vzorové riešenie

Stačí použiť bežný trik. Ak nevieme jednoducho spočítať počet dobrých ciest, možno vieme ľahko zistiť počet zlých. Dobré potom zrátame tak, že od počtu všetkých odpočítame počet zlých.

Už vieme, že počet všetkých postupností je \(n^k\).

Ako vyzerá zlá postupnosť? Neprejdeme počas nej po žiadnej označenej hrane. Skúsme teda zo stromu odstrániť všetky označené hrany. Týmto sa nám strom rozpadne na niekoľko komponentov. Každá zlá postupnosť potom musí byť celá v jednom komponente. Totiž, keby zahŕňala zlá postupnosť vrcholy z rôznych komponentov, museli by sme niekedy prejsť z jedného komponentu do druhého. Jedinou možnosťou, ako by sme to vedeli, je označená hrana. To je ale spor s predpokladom, že táto postupnosť je zlá.

Keď teda chceme zistiť počet zlých postupností, stačí nám ich vypočítať pre každý komponent zvlášť. Hodnoty pre jednotlivé komponenty potom iba sčítame.

Koľko zlých postupností dĺžky \(k\) je v komponente veľkosti \(x\)? To už predsa vieme, je to niečo podobné ako počet postupností v celom strome, akurát teraz nemá veľkosť \(n\), ale \(x\). Bude to teda \(x^k\).

Nech sa nám odobratím označených hrán strom rozpadol na \(c\) komponentov. Ich veľkosti sú \(x_1, \ldots, x_c\). Potom počet dobrých postupností je \(n^k - \sum_{i=1}^{c} x_i^k\).

Na zrátanie veľkostí komponentov vieme použiť jedno DFS.

Takéto riešenie by malo získať \(6\) bodov. Na plný počet musíme naviac použiť umocňovanie v logaritmickom čase. O tomto jednoduchom triku si môžete viac prečítať tu.

DFS potrebuje \(O(n)\) času. Následne veľkosť každého komponentu umocníme na \(k\)-tu. Komponentov je najviac \(n\). Takéto riešenie teda bude mať časovú zložitosť \(O(n \log k)\).

Pamätať si potrebujeme celý strom. Pamäťová zložitosť je preto \(O(n)\).

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

ll n, k, a, b, c;
vector<vector<ll>> G;
vector<bool> vis;

const ll MOD = 1000000007;

// Na umocnovanie si spravime vlastnu funkciu, lebo to potrebujeme modulovat
ll expo(ll a, ll x) {
    if(x == 0LL)
        return 1LL;
    if(x == 1LL)
        return a;
    ll res = expo(a, x / 2);
    res = (res * res) % MOD;
    if(x & 1LL)
        res = (res * a) % MOD;
    return res;
}

ll dfs(ll u) {
    int comp_sz = 1;
    vis[u] = true;
    for(int i = 0; i < G[u].size(); i++) {
        ll v = G[u][i];
        if(!vis[v])
            comp_sz += dfs(v);
    }
    return comp_sz;
}

int main() {
    cin >> n >> k;
    G.resize(n);
    vis.assign(n, false);
    for(int i = 0; i < n - 1; i++) {
        cin >> a >> b >> c; a--; b--;
        // Hrany s exkrementom rovno odignorujeme
        if(c == 0) {
            G[a].push_back(b);
            G[b].push_back(a);
        }
    }
    ll bad = 0;
    ll all = expo(n, k);
    for(int i = 0; i < n; i++)
        if(!vis[i])
            bad += expo(dfs(i), k);
    // Ak nechapete, preco takto modulujeme, pozrite si, ako funguje % v C++.
    cout << (((all - bad) % MOD) + MOD) % MOD << '\n';


    return 0;
}

import sys


sys.setrecursionlimit(200000)

MOD = 1000000007


def expo(a, x):
    if x == 0:
        return 1
    if x == 1:
        return a
    res = expo(a, x // 2)
    res = (res * res) % MOD
    if x & 1:
        res = (res * a) % MOD
    return res


def dfs(u, vis, G):
    comp_sz = 1
    vis[u] = True
    for v in G[u]:
        if not vis[v]:
            comp_sz += dfs(v, vis, G)
    return comp_sz


n, k = [int(x) for x in input().split()]
vis = [False] * n
G = [[] for _ in range(n)]
for i in range(n - 1):
    a, b, c = [int(x) for x in input().split()]
    a -= 1
    b -= 1
    if c == 0:
        G[a].append(b)
        G[b].append(a)

bad = 0
all_sequences = expo(n, k)
for i in range(n):
    if not vis[i]:
        bad += expo(dfs(i, vis, G), k)

print(((all_sequences - bad) % MOD + MOD) % MOD)