HackerRank Cut Tree problem solution

In this HackerRank Cut, Tree problem solution we have given a tree T with n nodes, how many subtrees (T’) of T have at most K edges connected to (T – T’)?

Problem solution in Python.

line = input().split()
n = int(line[0])
K = int(line[1])

adjlst = [[1]]
for i in range(n):
    adjlst.append([])
adjlst[1].append(0)

for i in range(n - 1):
    line = input().split()
    x = int(line[0])
    y = int(line[1])
    adjlst[x].append(y)
    adjlst[y].append(x)

def dfs(node, par, adjlst, K):
    conlst = [0] * (K + 1)
    dislst = [0] * (K + 1)
    conlst[0] = 1
    for chld in adjlst[node]:
        if chld != par:
            tcl, tdl = dfs(chld, node, adjlst, K)
            # print(str(tcl))
            # print(str(tdl))
            dislst[0] += tdl[0]
            tcl2 = [0] * (K + 1)
            for i in range(K):
                dislst[i + 1] += tcl[i] + tdl[i + 1]
                tcl2[i + 1] += conlst[i]
            for i in range(K + 1):
                for j in range(K + 1):
                    if i + j <= K:
                        tcl2[i + j] += tcl[i] * conlst[j]
            conlst = list(tcl2)
    return conlst, dislst

conlst, dislst = dfs(1, 0, adjlst, K)
# print(str(conlst))
# print(str(dislst))
ans = sum(conlst) + sum(dislst) + 1
print(str(ans))

Problem solution in Java.

// practice with kaiboy
import java.io.*;
import java.util.*;

public class cuttree extends PrintWriter {
    cuttree() { super(System.out, true); }
    Scanner sc = new Scanner(System.in);
    public static void main(String[] $) {
        cuttree o = new cuttree(); o.main(); o.flush();
    }

    int[] eo; int[][] ej;
    void append(int i, int j) {
        int o = eo[i]++;
        if (o >= 2 && (o & o - 1) == 0)
            ej[i] = Arrays.copyOf(ej[i], o << 1);
        ej[i][o] = j;
    }
    long[][] dp; long[] tt; int k;
    void init(int n) {
        eo = new int[n]; ej = new int[n][2];
        dp = new long[n][k + 1]; tt = new long[k + 1];
    }
    void mult(long[] aa, long[] bb) {
        Arrays.fill(tt, 0);
        for (int i = 0; i <= k; i++)
            for (int j = 0; i + j <= k; j++)
                tt[i + j] += aa[i] * bb[j];
        System.arraycopy(tt, 0, aa, 0, k + 1);
    }
    long ans = 1;
    void dfs(int p, int i) {
        dp[i][0] = 1;
        for (int o = eo[i]; o-- > 0; ) {
            int j = ej[i][o];
            if (j != p) {
                dfs(i, j);
                dp[j][1]++;
                mult(dp[i], dp[j]);
            }
        }
        for (int h = p == -1 ? k : k - 1; h >= 0; h--)
            ans += dp[i][h];
    }
    void main() {
        int n = sc.nextInt();
        if (n == 1) {
            println(2);
            return;
        }
        k = sc.nextInt();
        if (k == n)
            k = n - 1;
        init(n);
        for (int h = 0; h < n - 1; h++) {
            int i = sc.nextInt() - 1;
            int j = sc.nextInt() - 1;
            append(i, j);
            append(j, i);
        }
        dfs(-1, 0);
        println(ans);
    }
}

Problem solution in C++.

#include <cmath>
#include <cstdio>
#include <vector>
#include <iostream>
#include <algorithm>
#include<climits>
#include <numeric>
using namespace std;

class Tree {
  int n_;
  std::vector<std::vector<int>> edges;
  std::vector<int64_t> cut;
  std::vector<std::vector<int64_t>> connect;

  void dfs(int u, int p, int k) {
    connect[u][0] = 1;
    for (auto& v : edges[u]) {
      if (v == p) {
        continue;
      }
      dfs(v, u, k);
      for (int ku = k; ku >= 0; --ku) {
        int64_t sum = 0;
        for (int kv = 0; kv <= ku; ++kv) {
          sum += connect[u][ku-kv] * (connect[v][kv] + (kv == 1 ? 1 : 0));
        }
        connect[u][ku] = sum;
      }
    }
    cut[u] = std::accumulate(connect[u].begin(), connect[u].begin() + k, 0ll);
  }
  
public:
  Tree(int n): n_(n), edges(n) {
  }
  void add(int u, int v) {
    edges[u].push_back(v);
    edges[v].push_back(u);
  }
  int64_t cutWays(int k) {
    cut.assign(n_, 0);
    connect.assign(n_, std::vector<int64_t>(k+1));
    dfs(0, 0, k);
    cut[0] += connect[0][k];
    // add an empty subtree!
    return 1 + std::accumulate(cut.begin(), cut.end(), 0ll);
  }
};

int main() {
    /* Enter your code here. Read input from STDIN. Print output to STDOUT */   
    for (int n, k, u, v; std::cin >> n >> k; ) {
    Tree tree(n);
    for (int i = 0; i < n - 1; ++i) {
      std::cin >> u >> v;
      tree.add(u-1, v-1);
    }
    std::cout << tree.cutWays(k) << std::endl;
  }
    
    return 0;
}

Problem solution in C.

#include <stdio.h>
#include <stdlib.h>
void dfs(int x,int parent);
int table[50][50]={0},c[50]={0},p[50],q[51]={0},N;
long long dp1[50][51]={0},dp2[50][51]={0},temp[51],ans; // 1 without 2 with [node][K]
int main(){
  int K,x,y,i,j,k;
  scanf("%d%d",&N,&K);
  for(i=0;i<N-1;i++){
    scanf("%d%d",&x,&y);
    table[x-1][y-1]=-1;
    table[y-1][x-1]=-1;
  }
  dfs(0,0);
  for(i=0;i<N;i++)
    if(!c[i])
      q[++q[0]]=i;
  while(q[0]){
    x=q[q[0]--];
    dp1[x][0]=dp2[x][0]=1;
    for(i=0;i<N;i++)
      if(table[x][i]==1){
        for(j=1;j<=K;j++)
          temp[j]=0;
        for(j=1;j<=K;j++){
          dp1[x][j]+=dp1[i][j]+dp2[i][j-1];
          for(k=0;k<=j;k++)
            if(k==1)
              temp[j]+=(dp2[i][k]+1)*dp2[x][j-k];
            else
              temp[j]+=dp2[i][k]*dp2[x][j-k];
        }
        for(j=1;j<=K;j++)
          dp2[x][j]=temp[j];
      }
    if(x){
      c[p[x]]--;
      if(!c[p[x]])
        q[++q[0]]=p[x];
    }
  }
  for(i=0,ans=0;i<=K;i++)
    ans+=dp1[0][i]+dp2[0][i];
  printf("%lld",ans);
  return 0;
}
void dfs(int x,int parent){
  int i;
  p[x]=parent;
  for(i=0;i<N;i++)
    if(table[x][i] && i!=parent){
      table[x][i]+=2;
      c[x]++;
      dfs(i,x);
    }
  return;
}