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;
}