Codeforces Round #321 (Div. 2) C. Kefa and Park

原题: http://codeforces.com/contest/580/problem/C


题目大意: 给一个tree, tree上有几个mark的节点, 连续走m个mark的节点就不能继续往下走了. 问你能走到多少个leaf. leaf的定义是没有child的节点.


分析: 首先找出mark的节点的index, 存一下, 然后找到leaf的index, 存一下, 找leaf时, 看一下相邻的节点个数就可以判断出那个是leaf, 如果节点个数是1(只有父节点), 那么就是leaf. 然后从root dfs一下, 每次搜到一个mark的节点就计数, 达到m个就return, 然后遍历每个leaf节点, 如果这些节点都visited了, 那么就res++.

  int res = 0;
    ArrayList<Integer>[] g;
    ArrayList<Integer> leaf;
    int[] cat;
    int n;
    int m;
    boolean[] visited;
    public void solve(int testNumber, InputReader in, OutputWriter out) {
        n = in.readInt();
        m = in.readInt();
        g = new ArrayList[n];
        visited = new boolean[n];
        leaf = new ArrayList<>();
        cat = IOUtils.readIntArray(in, n);
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>();
        }
        for (int i = 1; i < n; i++) {
            int x = in.readInt();
            int y = in.readInt();
            x--;
            y--;
            g[x].add(y);
            g[y].add(x);
        }
        for (int i = 1; i < n; i++)
            if (g[i].size() == 1)
                leaf.add(i);
        dfs(0,0);
        for(Integer i:leaf)
            if(visited[i])
                res++;
        out.print(res);
    }

    private void dfs(int v, int c) {
        if (cat[v] == 1)
            c++;
        else
            c = 0;
        if (c > m)
            return;
        visited[v] = true;
        for (int i :g[v]) {
            if (!visited[i])
                dfs(i, c);
        }
        return;
    }