// Computes the maximum number of queens with C colours that you can place on a NxN grid such that
// queens of different colour do not attack each other.
//
// This program can be used to compute the best known solutions for A250000, as well as solutions for the 
// derivative sequences such as A308632 and A328283.
//
// The program uses a version of hill climbing. The key to the approach is a very fast mutation operator.
// Mutations involve changing the value of a single cell: empty to a queen, queen to an empty or
// queen of colour C1 to colour C2. Checking whether such a mutation improves the score takes time O(C),
// while making the actual change and updating auxiliary variables takes time O(N).
// 
// The scoreType can play an important role and produce different results.
// I found that the "extra" options seems to be the best, but you need to play around with it.
// You can also change the number of random starting solutions Q in solveMany().
//
// Compilation: javac A250000.java
// Usage: java A250000 N C
// N is the grid size
// C is the number of colours
//
// Dmitry Kamenetsky, 17th of October 2019.


import java.util.*;

public class A250000
{
  public static int N;    //grid size
  public static int C;    //number of colours
  
  final static int MULT=1000;
  
  //basic - maximize minQueens
  //total - maximize minQueens then totalQueens
  //extra - maximize minQueens and minimize extra queens needed to add
  final static String scoreType="extra";
  
  public static void main(String[] args)
  {
    if (args.length!=2)
    {
      System.out.println("Usage: java A250000 N C");
      System.exit(0);
    }
    
    N=Integer.parseInt(args[0]);
    C=Integer.parseInt(args[1]);
    
    solveMany();
  }
   
  
  public static void solveMany()
  {
    int maxChanges=5;
    int Q=20;
    int bestScore=Integer.MIN_VALUE;    
    int[] bestScores=new int[Q];
    Arrays.fill(bestScores,Integer.MIN_VALUE);
    
    int[][][] bestA=new int[Q][N][N];
    int[][] a=new int[N][N];
    
    while(true)
    {
      for (int q=0; q<Q; q++)
      {
        for (int r=0; r<N; r++) for (int c=0; c<N; c++) a[r][c]=bestA[q][r][c];
        
        int changes=(int)(Math.random()*maxChanges+1);
        for (int i=0; i<changes; i++)
        {
          int r=(int)(Math.random()*N);
          int c=(int)(Math.random()*N);
          int val=(int)(Math.random()*(C+1));
          a[r][c]=val;          
        }
        
        optimizeChangesFast(a);
        
        int score=score(a);
        
        if (score>=bestScores[q])
        {
          bestScores[q]=score;          
          for (int r=0; r<N; r++) for (int c=0; c<N; c++) bestA[q][r][c]=a[r][c];
        }
        
        if (score>bestScore)
        {
          System.out.println("score "+score);
          printNice(a);          
          bestScore=score;
        }        
      }
    }
  }  


  //O(C) check, O(N) update
  public static void optimizeChangesFast(int[][] a)
  {
    int[][][] counts=new int[C+1][N][N];   //number of times this location is covered by a queen of colour c
    
    int bad=0;
    int[] queens=new int[C+1];    
                 
              
    for (int r=0; r<N; r++)
      for (int c=0; c<N; c++)
      {
        if (a[r][c]==0) continue;   //empty cell
        
        queens[a[r][c]]++;
        
        for (int r2=0; r2<N; r2++)
          for (int c2=0; c2<N; c2++)
            //queen at (r,c) attacks location at (r2,c2)
            if (r2==r || c2==c || Math.abs(r2-r)==Math.abs(c2-c))
            {
              counts[a[r][c]][r2][c2]++;
              //queen at (r,c) attacks queen at (r2,c2)
              if (a[r2][c2]>0 && a[r2][c2]!=a[r][c]) bad++;
            }
      }    

      
    int score=scoreHelper(queens) - bad*MULT;

    //if (score!=score(a)) System.out.println("aaargh2 "+score+" "+score(a));
    //if (1==1) return;
    
    
    int C2=C+1;
    int[] ind=new int[N*N*C2];
    for (int i=0; i<ind.length; i++) ind[i]=i;
    
    while(true)
    {
      shuffle(ind);
      boolean changed=false;
      
      for (int i=0; i<ind.length; i++)
      {
        int r=ind[i]/(N*C2);
        int c=(ind[i]-r*N*C2)/C2;
        int val=ind[i]%C2;
        if (a[r][c]==val) continue;   //no change
                
        int old=a[r][c];        
        int newBad=bad;
        
        //added queen
        if (old==0 && val>0)
        {
          //if (1==1) return;
          queens[val]++;

          for (int k=1; k<=C; k++)
            if (k!=val) newBad+=2*counts[k][r][c];
        }
        //removed queen
        else if (old>0 && val==0)
        {
          //if (1==1) return;
          queens[old]--;
          
          for (int k=1; k<=C; k++)
            if (k!=old) newBad-=2*counts[k][r][c];          
        }
        //changed queen colours
        else
        {
          //if (1==1) return;
          queens[old]--;
          queens[val]++;

          for (int k=1; k<=C; k++)
          {
            if (k!=old) newBad-=2*counts[k][r][c];            
            if (k!=val) newBad+=2*counts[k][r][c];
          }    

          newBad-=2;    //to handle double counting         
        }
        
        //make move
        a[r][c]=val;        
       
        int newScore=scoreHelper(queens) - newBad*MULT;      
        
        //if (newScore!=score(a)) System.out.println("aaargh "+newScore+" "+score(a));
        //if (1==1) return;
        
        
        if (newScore>=score)
        {
          if (newScore>score) changed=true;
          score=newScore;
          bad=newBad;

          //update counts 
          //horizontal
          for (int c2=0; c2<N; c2++)
          {
            if (c2==c) continue;
            if (old!=0) counts[old][r][c2]--;    
            if (val!=0) counts[val][r][c2]++;            
          }

          //vertical
          for (int r2=0; r2<N; r2++)
          {
            if (r2==r) continue;
            if (old!=0) counts[old][r2][c]--;    
            if (val!=0) counts[val][r2][c]++;            
          }

          // \
          int min=Math.min(r,c);
          for (int r2=r-min,c2=c-min; r2<N && c2<N; r2++,c2++)
          {
            if (r2==r && c2==c) continue;
            if (old!=0) counts[old][r2][c2]--;    
            if (val!=0) counts[val][r2][c2]++;            
          }
        
          // /
          min=Math.min(r,N-1-c);
          for (int r2=r-min,c2=c+min; r2<N && c2>=0; r2++,c2--)
          {
            if (r2==r && c2==c) continue;
            if (old!=0) counts[old][r2][c2]--;    
            if (val!=0) counts[val][r2][c2]++;            
          }   

          //same spot
          if (old!=0) counts[old][r][c]--;    
          if (val!=0) counts[val][r][c]++;                                                         
        }
        else
        {
          a[r][c]=old;
          
          //undo added queen
          if (old==0 && val>0)
            queens[val]--;
          //undo removed queen
          else if (old>0 && val==0)
            queens[old]++;
          //undo changed queen colours
          else
          {
            queens[old]++;
            queens[val]--;          
          }
        }
      }
      
      if (!changed) break;
    }      
  }  
  
  
  //computes score based on queen counts
  public static int scoreHelper(int[] queens)
  {
    int minQueens=Integer.MAX_VALUE;
    int totalQueens=0;
    for (int i=1; i<=C; i++)
    {
      minQueens=Math.min(minQueens,queens[i]);
      totalQueens+=queens[i];
    }  
    
    int extra=0;
    for (int i=1; i<=C; i++)
      if (queens[i]<minQueens+1) extra+=minQueens+1-queens[i];
    
   
    if (scoreType.equals("total"))
      return MULT*minQueens + totalQueens;
    else if (scoreType.equals("extra"))      
      return MULT*minQueens + (10-extra);
    else if (scoreType.equals("basic"))
      return minQueens;   
      
    return 0;
  }  
   

  //O(N^4)
  public static int score(int[][] a)
  {
    int bad=0;
    int[] counts=new int[C+1];
    
    for (int r=0; r<N; r++)
      for (int c=0; c<N; c++)
      {
        if (a[r][c]==0) continue;   //empty cell
        
        counts[a[r][c]]++;

        for (int r2=0; r2<N; r2++)
          for (int c2=0; c2<N; c2++)
            //queen at (r2,c2) attacks queen at (r,c)
            if (a[r2][c2]>0 && a[r2][c2]!=a[r][c] && (r2==r || c2==c || Math.abs(r2-r)==Math.abs(c2-c)))
              bad++;
      }
            
    return scoreHelper(counts) - bad*MULT;
  }
  
  
  //shuffle the array randomly
  public static void shuffle(int[] a)
  {
    for (int i=0; i<a.length; i++)
    {
      int k=(int)(Math.random()*(a.length-i)+i);
      int temp=a[i];  
      a[i]=a[k];
      a[k]=temp;  
    }   
  }      
  
  
  public static void printNice(int[][] a)
  {
    int[] queens=new int[C+1];    
    int totalQueens=0;
    
    for (int r=0; r<N; r++)
      for (int c=0; c<N; c++)
      {
        queens[a[r][c]]++;
        if (a[r][c]>0) totalQueens++;
      }
    
    System.out.println("N "+N+" C "+C+" total queens "+totalQueens);
    System.out.println("Queens by colour:");
    for (int k=1; k<=C; k++) System.out.println(k+": "+queens[k]);
    print(a);
    System.out.println();
  }
                 
  
  public static void print(int[][] a)
  {
    for (int i=0; i<N; i++)
    {
      for (int k=0; k<N; k++)
      {
        if (a[i][k]==0) System.out.print(".");
        else System.out.print(a[i][k]);
      }
      System.out.println();
    }
  }
}