Nastia and a Beautiful Matrix (Codeforces 1521E)
Tags: Greedy, binary search

You like numbers, don't you? Nastia has a lot of numbers and she wants to share them with you! Isn't it amazing?Let $$$a_i$$$ be how many numbers $$$i$$$ ($$$1 \le i \le k$$$) you have.An $$$n \times n$$$ matrix is called beautiful if it contains all the numbers you have, and for each $$$2 \times 2$$$ submatrix of the original matrix is satisfied:   The number of occupied cells doesn't exceed $$$3$$$;  The numbers on each diagonal are distinct. Make a beautiful matrix of minimum size.

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10\,000$$$) — the number of test cases.The first line of each test case contains $$$2$$$ integers $$$m$$$ and $$$k$$$ ($$$1 \le m, k \le 10^5$$$) — how many numbers Nastia gave you and the length of the array $$$a$$$, respectively.The second line of each test case contains $$$k$$$ integers $$$a_1, a_2, \ldots, a_{k}$$$ ($$$0 \le a_i \le m$$$, $$$a_1 + a_2 + \ldots + a_{k} = m$$$), where $$$a_i$$$ is how many numbers $$$i$$$ you have.It's guaranteed that the sum of $$$m$$$ and $$$k$$$ in one test doesn't exceed $$$2 \cdot 10^5$$$.

Output: For each $$$t$$$ test case print a single integer $$$n$$$ — the size of the beautiful matrix.In the next $$$n$$$ lines print $$$n$$$ integers $$$b_{i, j}$$$ ($$$0 \le b_{i, j} \le k$$$; if position is empty, print $$$b_{i, j} = 0$$$) — the beautiful matrix $$$b$$$ you made up. 

Note: Note that $$$0$$$ in this problem represents a blank, not a number.Examples of possible answers for the first test case:$$$\begin{array}{cc} 1 & 1 \\ 4 & 0 \\ \end{array} \hspace{0,5cm} \begin{array}{cc} 1 & 4 \\ 1 & 0 \\ \end{array} \hspace{0,5cm} \begin{array}{cc} 4 & 0 \\ 1 & 1 \\ \end{array}$$$Examples of not beautiful matrices for the first test case:$$$\begin{array}{cc} 1 & 0 \\ 4 & 1 \\ \end{array} \hspace{0,5cm} \begin{array}{cc} 4 & 1 \\ 7 & 1 \\ \end{array} \hspace{0,5cm} \begin{array}{cc} 1 & 0 \\ 4 & 0 \\ \end{array}$$$The example of the not beautiful matrix for the second test case:$$$\begin{array}{cc} 3 & 4 & 0 & 2 & 2 \\ 3 & 2 & 3 & 3 & 0 \\ 0 & 1 & 0 & 0 & 0 \\ 3 & 0 & 0 & 0 & 0 \\ 2 & 1 & 3 & 3 & 3 \\ \end{array}$$$Everything is okay, except the left-top submatrix contains $$$4$$$ numbers.

Sample Inputs:
2
3 4
2 0 0 1
15 4
2 4 8 1

Sample Outputs:
2
4 1
0 1
5
3 0 0 2 2
3 2 3 3 0
0 1 0 4 0
3 0 0 0 0
2 1 3 3 3

Solution (Python):
from math import sqrt, ceil

def sol():
    (m, k) = map(int, input().split())
    A = list(map(int, input().split()))
    A = sorted(list(enumerate(A)), reverse=True, key=lambda x: x[1])
    B = []
    C = []
    for elem in A:
        if len(B) > len(C):
            C += [elem[0] + 1] * elem[1]
        else:
            B += [elem[0] + 1] * elem[1]
    if len(C) > len(B):
        (C, B) = (B, C)
    even_n = ceil(max(sqrt(4 * m / 3), sqrt(2 * len(B))))
    if even_n % 2 == 1:
        even_n += 1
    odd_n = even_n - 1
    if odd_n ** 2 - (odd_n - 1) ** 2 / 4 >= m and odd_n * (odd_n + 1) / 2 >= len(B):
        n = odd_n
    else:
        n = even_n
    if n % 2 == 0:
        C += [0] * (3 * n ** 2 // 4 - len(B) - len(C))
    else:
        C += [0] * (n ** 2 - (n - 1) ** 2 // 4 - len(B) - len(C))
    if n % 2 == 0:
        border = n ** 2 // 4
    else:
        border = (n - 1) * (n + 1) // 4
    (B, D) = (B[:border], B[border:])
    D += C[border:]
    C = C[:border]
    B_cnt = 0
    C_cnt = 0
    D_cnt = 0
    print(n)
    for i in range(n):
        for j in range(n):
            if i % 2 == 1 and j % 2 == 1:
                print(0, end=' ')
            elif i % 2 == 1:
                print(B[B_cnt], end=' ')
                B_cnt += 1
            elif j % 2 == 1:
                print(C[C_cnt], end=' ')
                C_cnt += 1
            else:
                print(D[D_cnt], end=' ')
                D_cnt += 1
        print()
for iter in range(int(input())):
    sol()

Solution (Python):
from math import sqrt, ceil
import sys

def putin():
    return map(int, sys.stdin.readline().split())

def sol():
    (m, k) = putin()
    A = list(putin())
    A = sorted(list(enumerate(A)), reverse=True, key=lambda x: x[1])
    B = []
    C = []
    for elem in A:
        if len(B) > len(C):
            C += [elem[0] + 1] * elem[1]
        else:
            B += [elem[0] + 1] * elem[1]
    if len(C) > len(B):
        (C, B) = (B, C)
    even_n = ceil(max(sqrt(4 * m / 3), sqrt(2 * len(B))))
    if even_n % 2 == 1:
        even_n += 1
    odd_n = even_n - 1
    if odd_n ** 2 - (odd_n - 1) ** 2 / 4 >= m and odd_n * (odd_n + 1) / 2 >= len(B):
        n = odd_n
    else:
        n = even_n
    if n % 2 == 0:
        C += [0] * (3 * n ** 2 // 4 - len(B) - len(C))
    else:
        C += [0] * (n ** 2 - (n - 1) ** 2 // 4 - len(B) - len(C))
    if n % 2 == 0:
        border = n ** 2 // 4
    else:
        border = (n - 1) * (n + 1) // 4
    (B, D) = (B[:border], B[border:])
    D += C[border:]
    C = C[:border]
    B_cnt = 0
    C_cnt = 0
    D_cnt = 0
    print(n)
    for i in range(n):
        for j in range(n):
            if i % 2 == 1 and j % 2 == 1:
                print(0, end=' ')
            elif i % 2 == 1:
                print(B[B_cnt], end=' ')
                B_cnt += 1
            elif j % 2 == 1:
                print(C[C_cnt], end=' ')
                C_cnt += 1
            else:
                print(D[D_cnt], end=' ')
                D_cnt += 1
        print()
for iter in range(int(sys.stdin.readline())):
    sol()

Solution (Python):
from math import sqrt, ceil
for iter in range(int(input())):
    (m, k) = map(int, input().split())
    A = list(map(int, input().split()))
    A = sorted(list(enumerate(A)), reverse=True, key=lambda x: x[1])
    B = []
    C = []
    for elem in A:
        if len(B) > len(C):
            C += [elem[0] + 1] * elem[1]
        else:
            B += [elem[0] + 1] * elem[1]
    if len(C) > len(B):
        (C, B) = (B, C)
    even_n = ceil(max(sqrt(4 * m / 3), sqrt(2 * len(B))))
    if even_n % 2 == 1:
        even_n += 1
    odd_n = even_n - 1
    if odd_n ** 2 - (odd_n - 1) ** 2 / 4 >= m and odd_n * (odd_n + 1) / 2 >= len(B):
        n = odd_n
    else:
        n = even_n
    if n % 2 == 0:
        C += [0] * (3 * n ** 2 // 4 - len(B) - len(C))
    else:
        C += [0] * (n ** 2 - (n - 1) ** 2 // 4 - len(B) - len(C))
    if n % 2 == 0:
        border = n ** 2 // 4
    else:
        border = (n - 1) * (n + 1) // 4
    (B, D) = (B[:border], B[border:])
    D += C[border:]
    C = C[:border]
    B_cnt = 0
    C_cnt = 0
    D_cnt = 0
    print(n)
    for i in range(n):
        for j in range(n):
            if i % 2 == 1 and j % 2 == 1:
                print(0, end=' ')
            elif i % 2 == 1:
                print(B[B_cnt], end=' ')
                B_cnt += 1
            elif j % 2 == 1:
                print(C[C_cnt], end=' ')
                C_cnt += 1
            else:
                print(D[D_cnt], end=' ')
                D_cnt += 1
        print()