HackerRank New Year Present problem solution

In this HackerRank New Year Present problem solution, Nina received an odd New Year’s present from a student: a set of N unbreakable sticks. Each stick has a length, L, and the length of the ith stick is li-1. Deciding to turn the gift into a lesson, Nina asks her students the following:

How many ways can you build a square using exactly 6 of these unbreakable sticks?

Problem solution in Python.

#!/bin/python3
import collections
import math
import os
import random
import re
import sys

def choose(n,k):
    if k>n:
        return 0
    k = min(k, n-k)
    num,den = 1,1 
    for i in range(k):
        num *= (n-i)
        den *= i+1
    return num//den

def squareCount(l):
    counter = collections.Counter(l)
    l = tuple(counter.keys())
    choose2 = collections.Counter()
    choose3 = collections.Counter()
    choose4 = collections.Counter()
    le = len(l)
    for n in counter:
        count = counter[n]
        if count >= 2:
            choose2[n] = choose(count,2)
            if count >= 3:
                choose3[n] = choose(count,3)
                if count>=4:
                    choose4[n] = choose(count,4)
    t1 = 0
    t2 = 0
    for i in range(le):
        if counter[2*l[i]] >= 2 and counter[l[i]] >= 4:
            t1 += choose4[l[i]]*choose2[2*l[i]]
        if counter[l[i]]>=2:
            for j in range(i+1,le):
                if counter[l[j]]>=2 and counter[l[i]+l[j]] >= 2:
                    t2 += choose2[l[i]]*choose2[l[j]]*choose2[l[i]+l[j]]
                    
    doubles = collections.Counter()
    for i in range(le):
        if counter[l[i]] >= 2 and counter[2*l[i]] >= 2:
            doubles[2*l[i]] = choose2[l[i]]
 
    pairs = collections.Counter()
    mpairs = collections.Counter()
    for i in range(le):
        for j in range(i+1,le):
            if counter[l[i]+l[j]] >= 2:
                pairs[l[i]+l[j]] += counter[l[i]]*counter[l[j]]
                mpairs[l[i]+l[j]] += counter[l[i]]**2*counter[l[j]]**2

    t3 = sum(pairs[s]*doubles[s]*choose2[s] for s in doubles if s in pairs)
    t4 = sum((pairs[s]**2 - mpairs[s])*choose2[s] for s in pairs)//2
    CD = collections.Counter()

    for d in range(le):
        if counter[l[d]]>=3:
            for c in range(le):
                if l[c] < l[d]:
                    CD[l[d]-l[c]] += choose3[l[d]]*counter[l[c]]

    s1 = 0
    s2 = 0
    s4 = 0
    
    for a in range(le):
        for b in range(a+1,le):
            s1 += 2*CD[l[a]+l[b]]*counter[l[a]]*counter[l[b]]
            if counter[l[a] + 2*l[b]] >= 3:
                s2 += 2*choose3[l[a] + 2*l[b]]*counter[l[a]]*counter[l[b]]**2
            if counter[2*l[a] + l[b]] >= 3:
                s2 += 2*choose3[2*l[a] + l[b]]*counter[l[b]]*counter[l[a]]**2
            if counter[l[b]] >= 2 and counter[l[a] + 2*l[b]] >= 3:
                s4 += counter[l[a]]*choose2[l[b]]*choose3[l[a]+2*l[b]]
            if counter[l[a]] >= 2 and counter[2*l[a] + l[b]] >= 3:
                s4 += counter[l[b]]*choose2[l[a]]*choose3[2*l[a]+l[b]]

    s = (s1 - s2)//6
    s3 = 0
    for a in range(le):
        if counter[l[a]] >= 3 and counter[3*l[a]]>=3:
            s3 += choose3[l[a]]*choose3[3*l[a]]
            
            
    return t1 + t2 + t3 + t4 + s + s3 + s4


if __name__ == '__main__':
    n = int(input())

    l = list(map(int, input().rstrip().split()))

    print(squareCount(l))
    

Problem solution in Java.

 
import java.util.Scanner;

public class Solution {

    public static void main(String[] args) {
        Scanner s = new Scanner(System.in);
        int N = s.nextInt();
        int R = 10000001;
        int[] lenM = new int[R];

        for (int i = 0; i < N; i++) {
            int it = s.nextInt();
            lenM[it]++;
        }
        Node[] nodes = new Node[R];
        Node[] nodes1 = new Node[R];
        int nodeLen = 0;
        int[]vi=new int[R];
        for (int i = 0; i < lenM.length; i++) {
            if (lenM[i] >= 1) {
                Node node = new Node(i, lenM[i]);
                vi[i]=nodeLen;
                nodes[nodeLen++] = node;
                nodes1[i] = node;
                
            }
        }
        int last=0,v;
        for(int i=R-1;i>=0;i--) {
            v=vi[i];
            if(v==0) {
                vi[i]=last;
            }else {
                last=vi[i];
            }
        }
        int max = nodes[nodeLen - 1].v;
        long[] cnTwo = new long[R];
        long[] cnThree = new long[R];
        for (int i = 2; i < R; i++) {
            cnTwo[i] = (long) i * (i - 1) / 2;
            cnThree[i] = cnTwo[i] * (i - 2) / 3;
        }
        int[] td = new int[R];

        long ways = 0;
        for (int i = 0; i < nodeLen; i++) {
            Node n1 = nodes[i];
            Node n2, n3, n4;
            int ind;
            if (n1.num >= 3) {
                ind = n1.v * 3;
                if (ind <= max) {
                    n3 = nodes1[ind];
                    if (n3 != null) {
                        ways += (long) cnThree[n1.num] * cnThree[n3.num];
                    }
                }
            }

            for (int j = i + 1; j < nodeLen - 1; j++) {
                n2 = nodes[j];
                ind = n1.v * 2 + n2.v;
                if (ind <= max) {
                    n3 = nodes1[ind];
                    if (n3 != null) {
                        ways += (long) cnTwo[n1.num] * n2.num * cnThree[n3.num];
                    }
                }
                ind = n1.v + n2.v * 2;
                if (ind <= max) {
                    n3 = nodes1[ind];
                    if (n3 != null) {
                        ways += (long) n1.num * cnTwo[n2.num] * cnThree[n3.num];
                    }
                }

                ind = n1.v + n2.v;
                if (ind <= max) {
                    n3 = nodes1[ind];
                    if (n3 != null) {
                        ways += (long) td[ind] * n1.num * n2.num * cnTwo[n3.num];
                        ways += (long) cnTwo[n1.num] * cnTwo[n2.num] * cnTwo[n3.num];
                        td[ind] += n1.num * n2.num;
                    }
                } else {
                    break;
                }

                int k = j + 1;
                n3 = nodes[k];

                ind = n1.v + n2.v + n3.v;
                if(ind>max)
                    continue;
                int L = vi[ind];
                for (; L < nodeLen; L++) {
                    n4 = nodes[L];
                    int idx3 = n4.v - n1.v - n2.v;
                    n3 = nodes1[idx3];
                    if (n3 != null) {
                        ways += (long) n1.num * n2.num * n3.num * cnThree[n4.num];
                    }
                }

            }

        }

        for (int i = 0; i < nodeLen; i++) {
            Node n1 = nodes[i];
            Node n3;
            int idx = n1.v * 2;
            if (idx <= max) {
                n3 = nodes1[idx];
                if (n3 != null) {
                    ways += (long) td[n3.v] * cnTwo[n1.num] * cnTwo[n3.num];
                    if (n1.num >= 4)
                        ways += (long) n1.num * (n1.num - 1) * (n1.num - 2) * (n1.num - 3) / 24 * cnTwo[n3.num];
                }

            } else {
                break;
            }

        }
        System.out.println(ways);
        s.close();
    }

    public static class Node implements Comparable<Node> {
        private int v;
        private int num; 

        public Node(int v, int num) {
            this.v = v;
            this.num = num; 
        }

        public Node(int v) {
            this.v = v;
        }

        @Override
        public int compareTo(Node o) {

            return this.v < o.v ? -1 : (this.v == o.v ? 0 : 1);
        }

    }
}

Problem solution in C++.

#include <bits/stdc++.h>

using namespace std;

#define N 3030
#define M 10000001

typedef pair<int, int> pii;
vector <pii> vv;
vector <int> v;

int c[M];
int n, a[N];
long long ans;

void run1() {
    for(int i = 1; i <= n; i ++) c[a[i]] ++;
    for(int i = 1; i <= n; i ++) if(i >= 6 && a[i] != a[i + 1]) {
        int cnt = 0, j;
        for(j = i; a[j] == a[i]; j --) cnt ++;
        if(cnt < 2) continue;
        int tp = cnt * (cnt - 1) / 2;
        vv.clear(); v.clear();
        int u = 0;
        for(; j; j --) if(a[j] != a[j-1]){
            if(a[j] * 2 < a[i]) break;
            if(a[j] * 2 == a[i]) {
                u = c[a[j]];
            } else {
                int x = a[i] - a[j];
                vv.push_back(pii(c[a[j]], c[x]));
            }
        }
        ans += 1LL * tp * u * (u - 1) * (u - 2) * (u - 3) / 24;
        for(j = 0; j < (int)vv.size(); j ++) {
         ans += 1LL * tp * vv[j].first * (vv[j].first - 1) * (vv[j].second - 1) * vv[j].second / 4;
         v.push_back(vv[j].first * vv[j].second);
        }
        if(u > 1) v.push_back(u * (u - 1) / 2);
        if((int)v.size() > 1) {
            long long sum = 0;
            for(j = 0; j < (int)v.size(); j ++) sum += v[j];
            sum = sum * sum;
            for(j = 0; j < (int)v.size(); j ++) sum -= 1LL * v[j] * v[j];
            ans += sum * tp / 2;
        }
    }
    for(int i = 1; i <= n; i ++) c[a[i]] --;
}

void run2() {
    for(int i = 1; i < n; i ++) {
        if(i > 2) {
            for(int j = i + 1; j <= n; j ++) 
                           if(a[i] < a[j] && a[j] != a[j + 1]) {
                int cnt = 0;
                for(int k = j; a[k] == a[j]; k --) cnt ++;
                int x = a[j] - a[i];
                if(cnt > 2) {
                    ans += 1LL * c[x] * cnt * (cnt - 1) * (cnt - 2) / 6;
                }
            }
        }
        for(int j = 1; j < i; j ++) {
            int x = a[j] + a[i];
            if(x < M) c[x] ++;
        }
    }
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i ++) scanf("%d", a + i);
    sort(a + 1, a + n + 1);
    run1();
    run2();
    cout << ans << endl;
    return 0;
}

Problem solution in C.

#include <stdio.h>
#include <stdlib.h>
long long C(long long n,long long r);
void sort_a(int*a,int*c,int size,int*new_size);
void merge(int*a,int*left_a,int*right_a,int*c,int*left_c,int*right_c,
int left_size, int right_size,int*new_size);
int a[3000],c[3000],one[10000000]={0};
long long two[10000000]={0};

int main(){
  int N,M,i,j;
  long long ans=0,t1,t2,a2=0,a3=0;
  scanf("%d",&N);
  for(i=0;i<N;i++){
    scanf("%d",a+i);
    one[a[i]-1]++;
    c[i]=1;
  }
  for(i=0;i<N-1;i++)
    for(j=i+1;j<N;j++)
      if(a[i]+a[j]<=10000000)
        two[a[i]+a[j]-1]++;
  sort_a(a,c,N,&M);
  for(i=0;i<M;i++){
    if(c[i]>1){
      for(j=t1=t2=0;a[j]<=a[i]/2 && j<i;j++)
        if(a[j]*2==a[i]){
          if(c[j]>1)
            t2+=t1*C(c[j],2);
          if(c[j]>3)
            t2+=C(c[j],4);
        }
        else if(one[a[i]-a[j]-1]){
          t2+=t1*one[a[i]-a[j]-1]*c[j];
          if(c[j]>1 && one[a[i]-a[j]-1]>1)
            t2+=C(c[j],2)*C(one[a[i]-a[j]-1],2);
          t1+=one[a[i]-a[j]-1]*c[j];
        }
      ans+=t2*C(c[i],2);
      a2+=t2*C(c[i],2);
    }
    if(c[i]>2){
      for(j=t1=0;j<i;j++)
        if(two[a[i]-a[j]-1]){
          t2=two[a[i]-a[j]-1];
          if(a[j]*3==a[i]){
            if(c[j]>2)
              t1+=C(c[j],3)*6;
            if(c[j]>1)
              t2-=C(c[j],2);
          }
          else if(a[i]-2*a[j]>0 && one[a[i]-2*a[j]-1]){
            if(c[j]>1)
              t1+=C(c[j],2)*one[a[i]-2*a[j]-1]*3;
            t2-=c[j]*one[a[i]-2*a[j]-1];
          }
          if(a[j]*3!=a[i] && (a[i]-a[j])/2*2==a[i]-a[j] && one[(a[i]-a[j])/2-1]>1){
            t1+=c[j]*C(one[(a[i]-a[j])/2-1],2)*3;
            t2-=C(one[(a[i]-a[j])/2-1],2);
          }
          t1+=t2*c[j]*2;
        }
      ans+=t1*C(c[i],3)/6;
      a3+=t1*C(c[i],3)/6;
    }
  }
  printf("%lld",ans);
  //printf("%lld %lld %lld",ans,a2,a3);
  return 0;
}
long long C(long long n,long long r){
  int i;
  long long ans=1;
  for(i=0;i<r;i++)
    ans*=(n-i);
  for(i=2;i<=r;i++)
    ans/=i;
  return ans;
}
void sort_a(int*a,int*c,int size,int*new_size){
  if (size < 2){
    (*new_size)=size;
    return;
  }
  int m = (size+1)/2,i;
  int *left_a,*right_a,*left_c,*right_c;
  left_a=(int*)malloc(m*sizeof(int));
  right_a=(int*)malloc((size-m)*sizeof(int));
  left_c=(int*)malloc(m*sizeof(int));
  right_c=(int*)malloc((size-m)*sizeof(int));
  for(i=0;i<m;i++){
    left_a[i]=a[i];
    left_c[i]=c[i];
  }
  for(i=0;i<size-m;i++){
    right_a[i]=a[i+m];
    right_c[i]=c[i+m];
  }
  int new_l_size=0,new_r_size=0;
  sort_a(left_a,left_c,m,&new_l_size);
  sort_a(right_a,right_c,size-m,&new_r_size);
  merge(a,left_a,right_a,c,left_c,right_c,new_l_size,new_r_size,new_size);
  free(left_a);
  free(right_a);
  free(left_c);
  free(right_c);
  return;
}
void merge(int*a,int*left_a,int*right_a,int*c,int*left_c,int*right_c,
int left_size, int right_size,int*new_size){
  int i = 0, j = 0,index=0;
  while (i < left_size|| j < right_size) {
    if (i == left_size) {
      c[index] = right_c[j];
      a[index++] = right_a[j];
      j++;
    } else if (j == right_size) {
      c[index] = left_c[i];
      a[index++] = left_a[i];
      i++;
    } else if (left_a[i] <= right_a[j]) {
      c[index] = left_c[i];
      a[index++] = left_a[i];
      i++;
    } else {
      c[index] = right_c[j];
      a[index++] = right_a[j];
      j++;
    }
    if(index>1&&a[index-2]==a[index-1]){
      c[index-2]+=c[index-1];
      index--;
    }
  }
  (*new_size)=index;
  return;
}