HackerRank King and Four Sons problem solution

In this HackerRank King and Four Sons problem solution The King tasks you with finding the number of ways of selecting K detachments of battalions to capture K countries using the criterion that is given by the problem.

Problem solution in Python.

import math

MOD = int(1e9 + 7)
ans=[[1], [1, 1], [1, 2, 1], [1, 3, 3, 1], [1, 4, 6, 4, 1]]

def S_(n, k):
    def cnk(n, k):
        return int(math.factorial(n) // (math.factorial(k) * math.factorial(n - k)))
    return sum(cnk(n, i) for i in range(k, n + 1, 4)) % MOD

def S(n, k = 0):
    if n < 5:
        return sum(ans[n][k::4])
    r = pow(2, n - 2, MOD) - pow(2, n // 2, MOD)
    if n & 1:
        r = (r + pow(2, (n - 3) // 2, MOD) * sum(ans[3][(k - (n - 3)// 2) % 4::4])) % MOD
    else:
        r = (r + pow(2, (n - 3) // 2, MOD) * sum(ans[4][(k - (n - 3)// 2) % 4::4])) % MOD
    return int(r)

n, k = map(int, input().split())
a = list(map(int, input().split()))
det = [S(i) for i in a]
mem = [[float("+inf")] * (i + 1) for i in range(n + 1)]
mem[0][0] = 1
for i in range(1, n + 1):
    mem[i][1] = sum(det[:i]) % MOD
    mem[i][i] = (mem[i - 1][i - 1] * det[i - 1]) % MOD
for i in range(3, n + 1):
    for j in range(2, min(i, k + 1)):
        mem[i][j] = (mem[i - 1][j] + mem[i - 1][j - 1] * det[i - 1]) % MOD
print(mem[n][k])

Problem solution in Java.

import java.io.*;
import java.util.*;

public class Solution {
  private static InputReader in;
  private static PrintWriter out;
  public static long mod = 1000000007;
  public static long INV4 = pow(new Complex(4,0), mod-2).a;
  
  static class Complex {
    public long a,b;
    public Complex(long a, long b) {
      this.a = a;
      this.b = b;
      if (a < 0) a += mod;
     
    }
    
    public Complex multiply(Complex other) {
      return new Complex((a * other.a + mod - (b * other.b % mod)) % mod, (a * other.b + b * other.a) % mod);
    }
  }
  
  public static Complex pow(Complex a, long e) {
    Complex r = new Complex(1,0);
    while(e>0) {
      if ((e&1)==1) r = r.multiply(a);
      a = a.multiply(a);
      e >>= 1;
    }
    return r;
  }
  
  public static long nways(long x) {
    Complex s1 = pow(new Complex(1, 1), x);
    Complex s2 = pow(new Complex(1, mod - 1), x);
    Complex s3 = pow(new Complex(2, 0), x);
    long res = (s3.a + s1.a + s2.a) % mod;
    return res * INV4 % mod;
  }

  public static void main(String[] args) throws IOException {
    in = new InputReader(System.in);
    out = new PrintWriter(System.out, true);

    int n = in.nextInt();
    int k = in.nextInt();
    long[] dp = new long[k+1];
    dp[0] = 1;
    for (int i = 0; i < n; i++) {
      long m = nways(in.nextInt());
      long[] next = new long[k+1];
      System.arraycopy(dp, 0, next, 0, k+1);
      for (int j = 0; j < k; j++) {
        next[j+1] = (next[j+1] + dp[j] * m) % mod;
      }
      dp = next;
    }
    out.println(dp[k]);
    out.close();
    System.exit(0);
  }

  static class InputReader {
    public BufferedReader reader;
    public StringTokenizer tokenizer;

    public InputReader(InputStream stream) {
      reader = new BufferedReader(new InputStreamReader(stream), 32768);
      tokenizer = null;
    }

    public String next() {
      while (tokenizer == null || !tokenizer.hasMoreTokens()) {
        try {
          tokenizer = new StringTokenizer(reader.readLine());
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
      }
      return tokenizer.nextToken();
    }

    public int nextInt() {
      return Integer.parseInt(next());
    }
  }


}

Problem solution in C++.

#include<bits/stdc++.h>
using namespace std;
#define FOR(i,a,b) for(int i = (a); i <= (b); ++i)
#define FORD(i,a,b) for(int i = (a); i >= (b); --i)
#define RI(i,n) FOR(i,1,(n))
#define REP(i,n) FOR(i,0,(n)-1)
#define mini(a,b) a=min(a,b)
#define maxi(a,b) a=max(a,b)
#define mp make_pair
#define pb push_back
#define st first
#define nd second
#define sz(w) (int) w.size()
typedef vector<int> vi;
typedef long long ll;
typedef long double ld;
typedef pair<int,int> pii;
const int inf = 1e9 + 5;
const int nax = 1e6 + 5;
const int mod = 1e9 + 7;

pii mul(pii a, pii b) {
    ll c = (ll) a.st * b.st - (ll) a.nd * b.nd;
    ll d = (ll) a.st * b.nd + (ll) a.nd * b.st;
    c %= mod;
    d %= mod;
    if(c < 0) c += mod;
    if(d < 0) d += mod;
    return mp((int) c, (int) d);
}

pii pw(pii a, int k) {
    pii r = mp(1, 0);
    while(k) {
        if(k % 2) r = mul(r, a);
        a = mul(a, a);
        k /= 2;
    }
    return r;
}
int pw(int a, int k) {
    int r = 1;
    while(k) {
        if(k % 2) r = (ll) r * a % mod;
        a = (ll) a * a % mod;
        k /= 2;
    }
    return r;
}

int f(int n) {
    int r = pw(mp(1,mod-1), n).st + pw(mp(1,1), n).st;
    r %= mod;
    r += pw(2, n);
    r %= mod;
    r = (ll) r * pw(4, mod - 2) % mod;
    return r;
}

int dp[105];

int main() {
    int n, k;
    scanf("%d%d", &n, &k);
    dp[0] = 1;
    REP(_, n) {
        int a;
        scanf("%d", &a);
        a = f(a);
        FORD(j, k, 1)
            dp[j] = (dp[j] + (ll) dp[j-1] * a) % mod;
    }
    printf("%dn", dp[k]);
    return 0;
}

Problem solution in C.

#include<stdio.h>
#include<stdlib.h>
#define CLOCK   1000000007
#define MODINV4 250000002
long int c[10000], M[10001][101];
long int fast_exp(long int x, int n, int clock)
{
    if( n == 0 )
    {
        return 1;
    }
    long int y = 1;
    while( n > 1 )
    {
        if( n % 2 == 0 )
        {
            x = x * x % clock;
            n = n / 2;
        }
        else
        {
            y = x * y % clock;
            x = x * x % clock;
            n = ( n - 1 ) / 2;
        }
    }
    return x * y % clock;
}
void compute_combinations(int n, int *army)
{
    int vals[4] = { 2, 2, 0, -4 };
    for( int i = 0 ; i < n ; i++ )
    {
        int div = army[i] / 4;
        int rem = army[i] % 4;
        if( div % 2 == 0 )
        {
            c[i] = fast_exp(2, army[i], CLOCK) + fast_exp(4, div, CLOCK) * ( vals[rem] + CLOCK );
        }
        else
        {
            c[i] = fast_exp(2, army[i], CLOCK) - fast_exp(4, div, CLOCK) * ( vals[rem] - CLOCK );
        }
        c[i] %= CLOCK;
        c[i] = c[i] * MODINV4 % CLOCK;
    }
}
int king(int n, int k, int* army)
{
    compute_combinations(n, army);
    for( int i = 0 ; i <= n ; i++ )
    {
        M[i][0] = 1;
    }
    for( int i = 1 ; i <= n ; i++ )
    {
        for( int j = 1 ; j <= k ; j++ )
        {
            M[i][j] = c[i - 1] * M[i - 1][j - 1] + M[i - 1][j];
            M[i][j] %= CLOCK;
        }
    }
    return M[n][k];
}
int main()
{
    FILE *fptr = fopen(getenv("OUTPUT_PATH"), "w");
    int n, k;
    scanf("%d %d", &n, &k);
    int *army = malloc(n * sizeof(int));
    for( int i = 0 ; i < n ; i++ )
    {
        scanf("%d", &army[i]);
    }
    fprintf(fptr, "%dn", king(n, k, army));
    fclose(fptr);
    return 0;
}