[CF580C] Kefa and Park

[CF580C] Kefa and Park

Description

给定一棵 N 个节点的有根树（根为 1 号节点），点权只有 1 和 0 两种，求从根节点到叶子节点的路径中，有多少条路径满足：路径上最大连续点权为 1 的节点个数不超过 M。

Solution

DFS 的时候记录一下已经连续了多少个 1

作为一个参数，如果超过 M，就直接退出

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1e6 + 5;
int n, m, a[N], ans;
vector<int> g[N];

void dfs(int p, int from, int con)
{
    if (a[p])
        con++;
    else
        con = 0;
    if (con > m)
        return;
    int flag = 0;
    for (int q : g[p])
        if (q != from)
        {
            dfs(q, p, con);
            flag = 1;
        }
    if (flag == 0)
        ++ans;
}

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n >> m;
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1; i < n; i++)
    {
        int t1, t2;
        cin >> t1 >> t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }
    dfs(1, 0, 0);
    cout << ans << endl;
}