Coloring (Codeforces 1774B)
Tags: Greedy, constructive algorithms

Cirno_9baka has a paper tape with $$$n$$$ cells in a row on it. As he thinks that the blank paper tape is too dull, he wants to paint these cells with $$$m$$$ kinds of colors. For some aesthetic reasons, he thinks that the $$$i$$$-th color must be used exactly $$$a_i$$$ times, and for every $$$k$$$ consecutive cells, their colors have to be distinct. Help Cirno_9baka to figure out if there is such a way to paint the cells.

Input: The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10\,000$$$)  — the number of test cases. The description of test cases follows.The first line of each test case contains three integers $$$n$$$, $$$m$$$, $$$k$$$ ($$$1 \leq k \leq n \leq 10^9$$$, $$$1 \leq m \leq 10^5$$$, $$$m \leq n$$$). Here $$$n$$$ denotes the number of cells, $$$m$$$ denotes the number of colors, and $$$k$$$ means that for every $$$k$$$ consecutive cells, their colors have to be distinct.The second line of each test case contains $$$m$$$ integers $$$a_1, a_2, \cdots , a_m$$$ ($$$1 \leq a_i \leq n$$$)  — the numbers of times that colors have to be used. It's guaranteed that $$$a_1 + a_2 + \ldots + a_m = n$$$.It is guaranteed that the sum of $$$m$$$ over all test cases does not exceed $$$10^5$$$.

Output: For each test case, output "YES" if there is at least one possible coloring scheme; otherwise, output "NO".You may print each letter in any case (for example, "YES", "Yes", "yes", and "yEs" will all be recognized as positive answers).

Note: In the first test case, there is no way to color the cells satisfying all the conditions.In the second test case, we can color the cells as follows: $$$(1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6)$$$. For any $$$2$$$ consecutive cells, their colors are distinct.

Sample Inputs:
2
12 6 2
1 1 1 1 1 7
12 6 2
2 2 2 2 2 2

Sample Outputs:
NO
YES

Solution (Python):
import sys
input = sys.stdin.readline
for _ in range(int(input())):
    (n, m, k) = map(int, input().split())
    a = list(map(int, input().split()))
    if (max(a) - 1) * k + a.count(max(a)) > n:
        print('NO')
    else:
        print('YES')

Solution (Python):
def solve():
    (n, m, k) = read_ints()
    aseq = read_ints()
    d = (n + k - 1) // k
    for ai in aseq:
        if ai > d:
            return 'NO'
    if n % k == 0:
        return 'YES'
    cnt = sum((ai == d for ai in aseq))
    return 'YES' if cnt <= n % k else 'NO'

def main():
    t = int(input()) if 1 else 1
    if 1:
        output = []
        for _ in range(t):
            ans = solve()
            output.append(ans)
        print_lines(output)
    else:
        for _ in range(t):
            solve()

def input():
    return next(test).strip()

def read_ints():
    return [int(c) for c in input().split()]

def print_lines(lst):
    print('\n'.join(map(str, lst)))
if __name__ == '__main__':
    import sys
    from os import environ as env
    if 'COMPUTERNAME' in env and 'L2A6HRI' in env['COMPUTERNAME']:
        sys.stdout = open('out.txt', 'w')
        sys.stdin = open('in.txt', 'r')
    test = iter(sys.stdin.readlines())
    main()

Solution (Python):
import sys
import math

def input_parse_expression(input_lines, i):
    (n, m, k) = map(int, input_lines[2 * i + 1].split())
    a = map(int, input_lines[2 * i + 2].split())
    return (n, m, k, a)

def solve(test_case_input):
    (n, m, k, a) = test_case_input
    if m < k:
        return 'NO'
    max_a_allowed = n // k
    num_1_gt_allowed = n % k
    for ai in a:
        if ai - max_a_allowed > 1:
            return 'NO'
        if ai - max_a_allowed == 1:
            num_1_gt_allowed -= 1
        if num_1_gt_allowed < 0:
            return 'NO'
    return 'YES'

def main():
    input_str = sys.stdin.read()
    input_lines = input_str.strip().split('\n')
    n_cases = int(input_lines[0])
    for i in range(n_cases):
        test_case_input = input_parse_expression(input_lines, i)
        solution = solve(test_case_input)
        print(solution)
if __name__ == '__main__':
    main()