HackerRank Lovely Triplets problem solution

In this HackerRank Lovely Triplets problem solution we have given P different lovely triplets and the minimum distance between each pair of nodes in the triplets Q and we need to draw a special graph.

Problem solution in Python.

#!/bin/python3

def choose(n, r):
    return n * (n-1) * (n-2) // (r * (r-1) * (r-2))
def product(xs):
    y = 1
    for x in xs:
        y *= x
    return y
def generate_primes(n):
    p = [True] * (n + 1)
    p[0] = False
    p[1] = False
    for i in range(n+1):
        if p[i]:
            yield i
            for j in range(2*i,n+1,i):
                p[j] = False
primes = list(generate_primes(5000))
def factorize(n):
    qs = []
    for p in primes:
        if p*p > n:
            break
        while n % p == 0:
            qs.append(p)
            n //= p
    if n != 1:
        qs.append(n)
    return qs
def core_size(q):
    return ((q - 1) // 2) * 3
def select_factors(p, q, width, memo):
    if p in memo:
        return memo[p]
    qs = [p, 1, 1]
    qv = core_size(q) + sum(qs)
    for i in range(min(width, p)):
        ps = [1, 1, 1]
        for r in factorize(p - i):
            ps[ps.index(min(ps))] *= r
        pv = core_size(q) + sum(ps)
        if i:
            nps, npv = select_factors(p - product(ps), q, width=width, memo=memo)
            pv += npv
        if pv < qv:
            qs = ps
            qv = pv
    memo[p] = (tuple(qs), qv)
    return memo[p]
def make_core(v, e, q):
    if q % 2 == 1:
        xs, v = [v, v+1, v+2], v+3
        for i in range(3):
            e.append((xs[i], xs[(i+1)%3]))
    else:
        xs, v = [v, v, v], v+1
    for i in range(3):
        for j in range(q//2 - 1):
            e.append((xs[i], v))
            xs[i] = v
            v += 1
    return xs, v
def make_triplets(v, e, q, ps):
    xs, v = make_core(v, e, q)
    for i in range(3):
        for _ in range(ps[i]):
            e.append((xs[i], v))
            v += 1
    return v
def make_coalesced_core(v, e, l):
    for i in range(l):
        e.append((v, v + i+1))
    v += l+1
    return v
p, q = map(int,input().split())
v = 0
e = []
if q == 2:
    while p:
        l = 3
        while choose(l+1,3) <= p:
            l += 1
        p -= choose(l,3)
        v = make_coalesced_core(v, e, l)
else:
    while p:
        ps, _ = select_factors(p, q, width=500, memo={})
        v = make_triplets(v, e, q, ps)
        p -= product(ps)
print(v, len(e))
assert v <= 100
for a, b in e:
    print(a+1, b+1)

Problem solution in Java.

import java.io.*;
import java.util.*;
import java.text.*;
import java.math.*;
import java.util.regex.*;

public class LovelyTriplets {
    public static void main(String[] args) throws Exception {
        Scanner in = new Scanner(System.in);
        int cases = 1;//in.nextInt();
        for (int testcase = 0; testcase < cases; testcase++) {
            int n = in.nextInt();
            int m = in.nextInt();
            lovelyTriplets(n, m);
        }
    }

    public static void lovelyTriplets(int n, int q) {
        if (q == 2) {
            lovelyTripletsTwo(n);
            return;
        }
        if (q == 3) {
            lovelyTripletsThree(n);
            return;
        }
        int best = Integer.MAX_VALUE;
        int bestIndex = -1;
        for (int i = 1; i < n; i++) {
            int[] b1 = bestThree(i);
            int[] b2 = bestThree(n-i);
            int sum = b1[0] +b1[1] + b1[2] + b2[0] + b2[1] + b2[2];
            if (sum < best) {
                best = sum;
                bestIndex = i;
            }
        }
        int[] b1 = bestThree(bestIndex);
        int[] b2 = bestThree(n-bestIndex);
        int sum = 0;
        for (int i : b1) sum+=i;
        for (int i : b2) sum+=i;
        int nodes = sum+3*(q-2);
        System.out.println(nodes + " " + nodes);
        int node = 7;
        for (int i = 0; i < b1[0]; i++) {
            System.out.println("1 " + node); node++;
        }
        for (int i = 0; i < b1[1]; i++) {
            System.out.println("2 " + node); node++;
        }
        for (int i = 0; i < b1[2]; i++) {
            System.out.println("3 " + node); node++;
        }
        for (int i = 0; i < b2[0]; i++) {
            System.out.println("4 " + node); node++;
        }
        for (int i = 0; i < b2[1]; i++) {
            System.out.println("5 " + node); node++;
        }
        for (int i = 0; i < b2[2]; i++) {
            System.out.println("6 " + node); node++;
        }
        System.out.println("1 4");
        System.out.println("2 5");
        System.out.println("3 6");
        int[] lasts = new int[]{4,5,6};
        for (int  i = 5; i <= q; i++) {
            for (int j = 0; j < 3; j++) {
                System.out.println(lasts[j] + " " + node);
                lasts[j] = node;
                node++;
            }
        }
        System.out.println(lasts[0] + " 2");
        System.out.println(lasts[1] + " 3");
        System.out.println(lasts[2] + " 1");

    }

    static void lovelyTripletsThree(int n) {
        int best = Integer.MAX_VALUE;
        int bestIndex = -1;
        for (int i = 1; i < n; i++) {
            int[] b1 = bestThree(i);
            int[] b2 = bestThree(n-i);
            int sum = b1[0] +b1[1] + b1[2] + b2[0] + b2[1] + b2[2];
            if (sum < best) {
                best = sum;
                bestIndex = i;
            }
        }
        int[] b1 = bestThree(bestIndex);
        int[] b2 = bestThree(n-bestIndex);
        int sum = 0;
        for (int i : b1) sum+=i;
        for (int i : b2) sum+=i;
        int nodes = sum+6;
        System.out.println(nodes + " " + nodes);
        int node = 7;
        for (int i = 0; i < b1[0]; i++) {
            System.out.println("1 " + node); node++;
        }
        for (int i = 0; i < b1[1]; i++) {
            System.out.println("2 " + node); node++;
        }
        for (int i = 0; i < b1[2]; i++) {
            System.out.println("3 " + node); node++;
        }
        for (int i = 0; i < b2[0]; i++) {
            System.out.println("4 " + node); node++;
        }
        for (int i = 0; i < b2[1]; i++) {
            System.out.println("5 " + node); node++;
        }
        for (int i = 0; i < b2[2]; i++) {
            System.out.println("6 " + node); node++;
        }
        System.out.println("1 2");
        System.out.println("2 3");
        System.out.println("3 1");
        System.out.println("4 5");
        System.out.println("5 6");
        System.out.println("6 4");

    }

    static int[] bestThree(int n) {
        int cube = (int)Math.pow(n, 1.0/3.0);
        for (int i = cube+1; i >= 1; i--) {
            if ((n%i) == 0) {
                int[] bt = bestTwo(n/i);
                return new int[]{i, bt[0], bt[1]};
            }
        }
        return new int[]{n, 1, 1};
    }

    static int[] bestTwo(int n) {
        int square = (int)Math.sqrt(n);
        for (int i = square+1; i >= 1; i--) {
            if ((n%i) == 0) {
                return new int[]{i, n/i};
            }
        }
        return new int[]{n, 1};
    }

    public static void lovelyTripletsTwo(int p) {
        int[] chooses = new int[34];
        for (int i = 3; i < 34; i++) {
            int val = i*(i-1)*(i-2)/6;
            chooses[i] = val;
        }



        int total = 0;
        List<Integer> flowers = new ArrayList<Integer>();
        while (total < p) {
            for (int i = 33; i >= 3; i--) {
                if (total + chooses[i] <= p) {
                    total += chooses[i];
                    flowers.add(i);
                    i++;
                }
            }
        }
        int numFlowers = flowers.size();
        int numLeaves = 0;
        for (int flower : flowers) numLeaves += flower;
        System.out.println((numFlowers + numLeaves) + " " + numLeaves);

        int nodeCount = 0;
        for (int flower : flowers) {
            int root = nodeCount+1;
            for (int i = 0; i < flower; i++) {
                System.out.println(root + " " + (root+i+1));
            }
            nodeCount = root + flower;
        }
    }

}

Problem solution in C++.

#include <cassert>
#include <iostream>
#include <vector>
#include <limits>
#include <array>

using std::cin;
using std::cout;
using std::cerr;
using std::vector;


namespace {
  struct Edge {
    int node1, node2;
  };
}

using Edges = vector<Edge>;


static int maxNodeNumber(const Edges &edges)
{
  int result = -1;
  int n_edges = edges.size();

  for (int i=0; i!=n_edges; ++i) {
    result = std::max(result,edges[i].node1);
    result = std::max(result,edges[i].node2);
  }

  return result;
}


namespace {
  struct Graph {
    const Edges edges;
    const int n_nodes;

    Graph(const Edges &edges_arg,int n_arg)
    : edges(edges_arg),
      n_nodes(n_arg)
    {
    }

    Graph(const Edges &edges_arg)
    : Graph(edges_arg,maxNodeNumber(edges_arg)+1)
    {
    }
  };
}


static int factorial(int n)
{
  int result = 1;
  while (n>1) {
    result *= n;
    --n;
  }
  return result;
}


static int comb(int n,int k)
{
  if (k==3) {
    assert(n>=3);
    int result = 1;
    int stop = n-k;
    while (n>stop) {
      result *= n;
      --n;
    }
    return result/factorial(3);
  }

  return factorial(n)/(factorial(k)*factorial(n-k));
}


static int nStarGraphTriplets(int n_spokes)
{
  return comb(n_spokes,3);
}


static Edges starGraphEdges(int n_spokes)
{
  Edges edges;
  for (int i=0; i!=n_spokes; ++i) {
    edges.push_back(Edge{0,i+1});
  }
  return edges;
}


static void offsetEdges(Edges &edges,int n)
{
  int n_edges = edges.size();

  for (int i=0; i!=n_edges; ++i) {
    edges[i].node1 += n;
    edges[i].node2 += n;
  }
}


static void
  addOffsetEdgesTo(Edges &new_edges,int a_n_nodes,const Edges &b_edges)
{
  Edges additional_edges = b_edges;
  offsetEdges(additional_edges,a_n_nodes);
  new_edges.insert(
     new_edges.end(),additional_edges.begin(),additional_edges.end()
  );
}


typedef std::array<int,3> Triple;


namespace {
  struct BarbCombo : std::array<Triple,2> {
    BarbCombo(const Triple &arg0,const Triple &arg1)
    : std::array<Triple,2>{{arg0,arg1}}
    {
    }
  };
}


static BarbCombo makeBarbResults(const vector<int> &v)
{
  return {{v[0],v[1],v[2]},{v[3],v[4],v[5]}};
}


static int barb9NodeCount(const BarbCombo &barb_counts)
{
  int total = 0;
  int n_groups = barb_counts.size();

  int l = 4;

  for (int i=0; i!=n_groups; ++i) {
    if (barb_counts[i][0]!=0 && barb_counts[i][1]!=0 && barb_counts[i][2]!=0) {
      total += 3*l+barb_counts[i][0]+barb_counts[i][1]+barb_counts[i][2];
    }
  }

  return total;
}


static int barb9EdgeCount(const BarbCombo &barb_counts)
{
  return barb9NodeCount(barb_counts);
}


static bool hasLegalNodeCounts(const BarbCombo &barb_results)
{
  int n_nodes = barb9NodeCount(barb_results);
  int n_edges = barb9EdgeCount(barb_results);

  return (n_nodes<=100 && n_edges<=100);
}


static int barbTotal(const BarbCombo &results)
{
  int total = 0;
  for (int i=0, n_results=results.size(); i!=n_results; ++i) {
    total += results[i][0]*results[i][1]*results[i][2];
  }
  return total;
}


// Try to find a set of six numbers such that a*b*c + d*e*f == n
// Make sure that we can create a legal graph from whatever we return.
static BarbCombo findBarbCombo(int n)
{
  vector<int> v = {1,1,0,1,1,1};

  assert(n>=1);

  v[5] = n;

  for (;;) {
    for (;;) {
      BarbCombo results = makeBarbResults(v);

      if (!hasLegalNodeCounts(results)) {
        break;
      }

      int total = barbTotal(results);

      if (total==n) {
        return results;
      }

      if (total>n) {
        break;
      }

      ++v[5];
    }

    int i = 5;
    while (v[i]==1) {
      --i;
    }
    v[i] = 1;
    ++v[i-1];
  }
}


static Edges barbGraphEdges(int q,const Triple &barb_counts)
{
  if (barb_counts[0]*barb_counts[1]*barb_counts[2]==0) {
    return {};
  }

  assert(q>=3);

  Edges edges;
  int base = 0;

  if (q%2==0) {
    edges = {{0,1},{0,2},{0,3}};
    base = 1;
  }
  else {
    edges = {{0,1},{0,2},{1,2}};
    base = 0;
  }

  int iter = 0;

  while (q>=5+iter*2) {
    for (int i=0; i!=3; ++i) {
      edges.push_back(Edge{base+iter*3+i,base+iter*3+i+3});
    }
    ++iter;
  }

  int offset = base+iter*3;
  int v = offset+3;

  for (int i=0; i!=3; ++i) {
    for (int j=0; j!=barb_counts[i]; ++j) {
      edges.push_back(Edge{i+offset,v});
      ++v;
    }
  }

  return edges;
}



static Graph operator+(const Graph &a,const Graph &b)
{
  Edges new_edges = a.edges;
  addOffsetEdgesTo(new_edges,a.n_nodes,b.edges);
  return Graph(new_edges);
}


static Edges makeMultiStarGraphEdges(int p)
{
  Edges edges;
  int n_nodes = 0;

  while (p>0) {
    int n_spokes = 3;
    while (nStarGraphTriplets(n_spokes+1)<=p) {
      ++n_spokes;
    }
    addOffsetEdgesTo(edges,n_nodes,starGraphEdges(n_spokes));
    n_nodes = maxNodeNumber(edges)+1;
    p -= nStarGraphTriplets(n_spokes);
  }

  return edges;
}


static Graph createSpecialGraph(int p,int q)
{
  if (q==2) {
    return Graph(makeMultiStarGraphEdges(p));
  }

  BarbCombo barb_combo = findBarbCombo(p);
  Edges edges0 = barbGraphEdges(q,barb_combo[0]);
  Edges edges1 = barbGraphEdges(q,barb_combo[1]);

  return Graph(edges0) + Graph(edges1);
}


int main()
{
  int p,q; cin >> p >> q;
  Graph g = createSpecialGraph(p,q);
  int m = g.edges.size();
  int n = g.n_nodes;
  cout << n << " " << m << "n";
  for (int i=0; i!=m; ++i) {
    cout << g.edges[i].node1+1 << " " << g.edges[i].node2+1 << "n";
  }
}

Problem solution in C.

#include <assert.h>
#include <limits.h>
#include <math.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

char* readline();
char** split_string(char*);

int** makeGraph(int p, int q, int startv, int *result_count){
    if(p == 0){
        *result_count = 1;
        int** toreturn = malloc(sizeof(int*));
        toreturn[0] = malloc(2*sizeof(int));
        toreturn[0][0] = 0;
        toreturn[0][1] = 0;
        return toreturn;
    }
    if(q == 2){
        int numspokes = 3;
        while(numspokes*(numspokes - 1)*(numspokes - 2) <= 6*p){
            numspokes++;
        }
        numspokes--;

        int tempresult_count;
        int** graph = makeGraph(p - (numspokes*(numspokes - 1)*(numspokes - 2))/6, q, startv + 1 + numspokes, &tempresult_count);
        *result_count = tempresult_count + numspokes;
        graph = realloc(graph, (*result_count)*sizeof(int*));
        for(int i = 0; i < numspokes; i++){
            graph[tempresult_count + i] = malloc(2*sizeof(int));
            graph[tempresult_count + i][0] = startv + 1;
            graph[tempresult_count + i][1] = startv + i + 2;
        }
        graph[0][0] += numspokes + 1;
        graph[0][1] += numspokes;
        return graph;
    }
    else{
        int factor1[q - 2];
        int factor2[q - 2];
        int factor3[q - 2];
        int currp = p;
        int extranum = 3*(q - 2);
        for(int i = 0; i < q - 2; i++){
            if(currp == 0){
                factor1[i] = 0;
                factor2[i] = 0;
                factor3[i] = 0;
            }
            else{
                factor1[i] = 1;
                while(factor1[i]*factor1[i]*factor1[i] <= currp){
                    factor1[i]++;
                }
                factor1[i]--;
                factor2[i] = factor1[i];
                while(factor1[i]*factor2[i]*factor2[i] <= currp){
                    factor2[i]++;
                }
                factor2[i]--;
                factor3[i] = factor2[i];
                while(factor1[i]*factor2[i]*factor3[i] <= currp){
                    factor3[i]++;
                }
                factor3[i]--;
                currp -= factor1[i]*factor2[i]*factor3[i];
                extranum += factor1[i] + factor2[i] + factor3[i];
            }
        }

        int tempresult_count;
        int** graph = makeGraph(currp, q, startv + extranum, &tempresult_count);
        *result_count = tempresult_count + extranum;
        graph = realloc(graph, (*result_count)*sizeof(int*));
        graph[tempresult_count] = malloc(2*sizeof(int));
        graph[tempresult_count][0] = startv + 1;
        graph[tempresult_count][1] = startv + 3*(q - 2);
        for(int i = 0; i < 3*(q - 2) - 1; i++){
            graph[tempresult_count + i + 1] = malloc(2*sizeof(int));
            graph[tempresult_count + i + 1][0] = startv + i + 1;
            graph[tempresult_count + i + 1][1] = startv + i + 2;
        }
        int index = tempresult_count + 3*(q - 2);
        int numvertex = 3*(q - 2);
        for(int i = 0; i < q - 2; i++){
            for(int j = 0; j < factor1[i]; j++){
                graph[index] = malloc(2*sizeof(int));
                graph[index][0] = startv + i + 1;
                graph[index][1] = startv + numvertex + 1;
                index++;
                numvertex++;
            }
            for(int j = 0; j < factor2[i]; j++){
                graph[index] = malloc(2*sizeof(int));
                graph[index][0] = startv + (q - 2) + i + 1;
                graph[index][1] = startv + numvertex + 1;
                index++;
                numvertex++;
            }
            for(int j = 0; j < factor3[i]; j++){
                graph[index] = malloc(2*sizeof(int));
                graph[index][0] = startv + 2*(q - 2) + i + 1;
                graph[index][1] = startv + numvertex + 1;
                index++;
                numvertex++;
            }
        }

        graph[0][0] += extranum;
        graph[0][1] += extranum;
        return graph;
    }    
}

int main()
{
    char** PQ = split_string(readline());

    char* P_endptr;
    char* P_str = PQ[0];
    int P = strtol(P_str, &P_endptr, 10);

    if (P_endptr == P_str || *P_endptr != '') { exit(EXIT_FAILURE); }

    char* Q_endptr;
    char* Q_str = PQ[1];
    int Q = strtol(Q_str, &Q_endptr, 10);

    if (Q_endptr == Q_str || *Q_endptr != '') { exit(EXIT_FAILURE); }

    int result_count;
    int** result = makeGraph(P, Q, 0, &result_count);
    for(int i = 0; i < result_count; i++){
        printf("%d %dn", result[i][0], result[i][1]);
    }
}

char* readline() {
    size_t alloc_length = 1024;
    size_t data_length = 0;
    char* data = malloc(alloc_length);

    while (true) {
        char* cursor = data + data_length;
        char* line = fgets(cursor, alloc_length - data_length, stdin);

        if (!line) { break; }

        data_length += strlen(cursor);

        if (data_length < alloc_length - 1 || data[data_length - 1] == 'n') { break; }

        size_t new_length = alloc_length << 1;
        data = realloc(data, new_length);

        if (!data) { break; }

        alloc_length = new_length;
    }

    if (data[data_length - 1] == 'n') {
        data[data_length - 1] = '';
    }

    data = realloc(data, data_length);

    return data;
}

char** split_string(char* str) {
    char** splits = NULL;
    char* token = strtok(str, " ");

    int spaces = 0;

    while (token) {
        splits = realloc(splits, sizeof(char*) * ++spaces);
        if (!splits) {
            return splits;
        }

        splits[spaces - 1] = token;

        token = strtok(NULL, " ");
    }

    return splits;
}