HackerEarth Proving your intelligence to your girlfriend solution

In this HackerEarth Proving your intelligence to your girlfriend problem solution Ashi is interested in checking the algorithmic skills of her boyfriend. So, she sets up a special 2-D graph for him. The 2-D graph looks like a maze with N rows each containing N nodes.

To connect these nodes, Ashi uses the following algorithm:

1. She selects two integers k1, k2 and computes the k1th and k2th fibonacci number, fib(k1) and fib(k2) respectively.
2. She takes the first row and selects the first two nodes in that row and adds an edge of weight (fib(k1) + fib(k2))%MOD. For, the next pair of consecutive nodes, she adds an edge of weight (fib(k1+1) + fib(k2+1))%MOD, then for the next pair (fib(k1+2) + fib(k2+2))%MOD and so on till the end of the row.
3. For the remaining rows, she uses the same algorithm incrementing the fibonacci number the same way as earlier.
4. Then she selects two more integers k3 and k4 and computes fib(k3) and fib(k4) respectively.
5. She, like filling rows, fills the columns by starting from the first column, finish adding edges in that and then move on to the next column

She wants her boyfriend to select a subset of edges such that all the nodes are connected with each other by some path and the total sum of edges in the subset is minimum over all possible subsets of edges that maintain the connectivity of the graph.

#include <stdio.h>
#include <vector>
#include <queue>
#include <assert.h>

using namespace std;

#define MAX 1000010
#define MOD 1000000007
#define LIM 1000000000000000000LL

class Edge{
  public:
    int from, to, wt;
    
    Edge(){}
    Edge(int f, int t, int w){
        from = f, to = t;
        wt = w;
    }
};

class Compare{
  public:
    bool operator() (const Edge& e1, const Edge& e2) const{
        return e1.wt > e2.wt;
    }
};

class MinSpanningTree{
    priority_queue< Edge, vector<Edge>, Compare > q;
    bool visited[MAX];
    vector< vector<Edge> > graph;
    int N;
    
  public:
    vector<Edge> mst;
    long long weight;
    
    MinSpanningTree(){}
    MinSpanningTree(vector< vector<Edge> >& _g, int _n) : graph(_g), N(_n){
        for (int i=0; i<N; i++)
            visited[i] = false;
        while (not q.empty())
            q.pop();
        mst.clear();
        weight = 0;
    }
    
    void ComputeMST(){
        int i, v, count = 0;
        Edge e;
        
        visited[0] = true;
        for (i=0; i<graph[0].size(); i++)
            q.push(graph[0][i]);
        while (count < N-1){
            e = q.top();
            q.pop();
            v = e.to;
            if (visited[v])
                continue;
            
            mst.push_back(e);
            weight += e.wt;
            count++;
            visited[v] = true;
            for (i=0; i<graph[v].size(); i++)
                if (not visited[ graph[v][i].to ])
                    q.push(graph[v][i]);
        }
        
        long long temp = 0;
        for (i=0; i<mst.size(); i++)
            temp += mst[i].wt;
        assert(mst.size() == N-1);
        assert(temp == weight);
    }
    
    void PrintMST(){
        for (int i=0; i<mst.size(); i++)
            printf("Wt (%d,%d) = %dn", mst[i].from, mst[i].to, mst[i].wt);
    }
};

int Fib(long long n){
  int i, j, k;
  long long fib[2][2]={{1,1},{1,0}},ret[2][2]={{1,0},{0,1}},tmp[2][2]={{0,0},{0,0}};
  while(n){
    if(n&1) {
      for(i=0;i<2;i++)
        for(j=0;j<2;j++)
          tmp[i][j]=0;
      for(i=0;i<2;i++) for(j=0;j<2;j++) for(k=0;k<2;k++)
          tmp[i][j]=(tmp[i][j]+ret[i][k]*fib[k][j])%MOD;
      for(i=0;i<2;i++) for(j=0;j<2;j++) ret[i][j]=tmp[i][j];
    }
    for(i=0;i<2;i++)
        for(j=0;j<2;j++)
          tmp[i][j]=0;
    for(i=0;i<2;i++) for(j=0;j<2;j++) for(k=0;k<2;k++)
                 tmp[i][j]=(tmp[i][j]+fib[i][k]*fib[k][j])%MOD;
          for(i=0;i<2;i++) for(j=0;j<2;j++) fib[i][j]=tmp[i][j];
    n/=2;

  }
  return (ret[0][1])%MOD;
}

int main(){
    int T, N, M, i, j, u, v, wt;
    long long k1, k2, k3, k4;
    int fib1, fib2, fib3, fib4;
    int fib11, fib22, fib33, fib44;
    int tmp1, tmp2, tmp3, tmp4;
    
  T = 1;
    while (T--) {
        scanf("%d", &N);
    assert(N>=1 and N<=1000);
        vector< vector<Edge> > g(N*N);
/*  expects 0-based indexing of vertices    */

        scanf("%lld %lld %lld %lld", &k1, &k2, &k3, &k4);
    assert(k1>=1 and k1<=LIM);
    assert(k2>=1 and k2<=LIM);
    assert(k3>=1 and k3<=LIM);
    assert(k4>=1 and k4<=LIM);
        fib11 = Fib(k1-1);
        fib1 = Fib(k1);
        fib22 = Fib(k2-1);
        fib2 = Fib(k2);
        fib33 = Fib(k3-1);
        fib3 = Fib(k3);
        fib44 = Fib(k4-1);
        fib4 = Fib(k4);
        
    //filling rows
        for (i=0; i<N; i++){
            for (j=0; j<N-1; j++){
                u = i*N+j;
                v = u + 1;
                wt = (fib1 + fib2)%MOD;
                g[u].push_back(Edge(u,v,wt));
                g[v].push_back(Edge(v,u,wt));
                tmp1 = fib11, tmp2 = fib22;
                fib11 = fib1, fib22 = fib2;
                fib1 = (fib11+tmp1)%MOD;
                fib2 = (fib22+tmp2)%MOD;
            }
        }
    //filling cols
        for (i=0; i<N; i++){
            for (j=0; j<N-1; j++){
                u = j*N+i;
                v = u + N;
                wt = (fib3 + fib4)%MOD;
                g[u].push_back(Edge(u,v,wt));
                g[v].push_back(Edge(v,u,wt));
                tmp3 = fib33, tmp4 = fib44;
                fib33 = fib3, fib44 = fib4;
                fib3 = (fib33+tmp3)%MOD;
                fib4 = (fib44+tmp4)%MOD;
            }
        }
       MinSpanningTree mstree(g,N*N);
        mstree.ComputeMST();
        printf("%lldn", mstree.weight);
        //mstree.PrintMST();
    }
    
    return 0;
}