CF 1062734 - Раскраска стены

CF 1062734 - \u0420\u0430\u0441\u043a\u0440\u0430\u0441\u043a\u0430 \u0441\u0442\u0435\u043d\u044b

Rating: -
Tags: -
Solve time: 38s
Verified: yes

Solution

Problem Understanding

We have a rectangular wall divided into tiles. Some positions contain lamps and cannot be painted. Every other position must receive one of k colors.

The restriction is local: take any horizontal or vertical continuous piece of tiles that is bounded by either the edge of the wall or a lamp on both sides. Inside such a piece, every painted tile must have a different color. In other words, colors cannot repeat inside any uninterrupted row segment or column segment.

The task is to decide whether such a coloring exists. If it does, we need to print one valid coloring, keeping lamp positions as 0.

The dimensions of the wall are at most 100 × 100, and the total number of cells over all test cases is at most 100000. This rules out algorithms that repeatedly try all color assignments. The number of possible colorings is enormous, so we need a direct construction. The small grid size allows linear or near-linear scans over the cells.

The main edge cases come from segments created by lamps. A coloring that works on the whole rectangle may fail when a segment is longer than the number of available colors, and checking only the whole row or whole column is not enough.

For example:

There is one column segment of length 5, but only 3 colors exist. The correct output is:

NO

A careless implementation might only look at the number of rows or columns and assume a construction always exists.

Another example:

The longest horizontal segment has length 4 in the second row. Since only 3 colors are available, the correct output is:

NO

The lamps do not help because that entire row has no separator.

A valid small case is:

The segments have maximum length 3, so the answer is YES. A repeating diagonal pattern can color all tiles without duplicates inside any segment.

Approaches

A direct brute-force solution would treat every tile as a variable and try different colors until a valid assignment is found. The condition can be checked after each attempt by scanning every row and column segment. This is correct because it explores every possible coloring, but the number of assignments is k raised to the number of tiles. Even with a small grid, this grows far beyond what can be processed.

The useful observation comes from looking at what makes a segment impossible. A segment containing L tiles needs L different colors, because no two positions inside it may share a color. If the longest segment length is larger than k, no coloring can exist.

The reverse direction is the key. If every segment has length at most k, we can use a fixed repeating pattern. Assign each tile the color (row + column) mod k + 1. Moving one step horizontally increases the color by one, and moving one step vertically does the same. Any consecutive segment of length at most k receives distinct residues, so every horizontal and vertical segment satisfies the rule.

The lamps simply interrupt these sequences. They do not need colors and are skipped during construction.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(k^cells)	O(cells)	Too slow
Optimal	O(nm)	O(nm)	Accepted

Algorithm Walkthrough

Scan every row and find the length of each consecutive group of tiles. If any group has length greater than k, the wall cannot be colored, because that single group already requires more colors than available.
Scan every column in the same way. A vertical segment longer than k also makes the answer impossible.
If all segments are short enough, color every tile using (row + column) mod k + 1. Keep every lamp as 0.
Print the constructed grid.

The reason the final pattern works is that every movement to an adjacent tile changes the color by exactly one modulo k. Since every uninterrupted segment contains at most k tiles, no color value can appear twice before the sequence wraps around.

Why it works

If a segment has more than k tiles, it needs more than k distinct colors, so reporting NO is unavoidable. If every segment has length at most k, consider any horizontal or vertical segment. The assigned colors on it form a sequence of consecutive values modulo k. A sequence of at most k consecutive residues cannot contain a duplicate, so every segment satisfies the requirement. Since every possible restricted segment is either horizontal or vertical, the whole wall is valid.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    ans = []

    for _ in range(t):
        n, m, k = map(int, input().split())
        wall = [list(map(int, input().split())) for _ in range(n)]

        ok = True

        for i in range(n):
            cur = 0
            for j in range(m):
                if wall[i][j] == 1:
                    cur += 1
                    if cur > k:
                        ok = False
                else:
                    cur = 0

        for j in range(m):
            cur = 0
            for i in range(n):
                if wall[i][j] == 1:
                    cur += 1
                    if cur > k:
                        ok = False
                else:
                    cur = 0

        if not ok:
            ans.append("NO")
            continue

        ans.append("YES")
        for i in range(n):
            row = []
            for j in range(m):
                if wall[i][j] == 0:
                    row.append("0")
                else:
                    row.append(str((i + j) % k + 1))
            ans.append(" ".join(row))

    print("\n".join(ans))

if __name__ == "__main__":
    solve()

The input is first stored because the lamp positions are needed both while checking segment lengths and while producing the answer.

The row and column scans maintain cur, the length of the current uninterrupted tile segment. A lamp resets this counter because it separates two independent segments. The moment cur exceeds k, the construction is impossible.

The coloring phase uses zero-based coordinates. The expression (i + j) % k + 1 produces values from 1 to k, while adding 1 avoids using 0, which is reserved for lamps. There is no need for special handling of small k, because the feasibility check already guarantees that every segment is short enough.

Worked Examples

Consider:

The scan and coloring process is:

Step	Position	Action	Current check result
Row 0	`1 1 0 1`	Finds segments of lengths 2 and 1	Valid
Row 1	`1 1 1 0`	Finds segment of length 3	Valid
Column 0	`1 1`	Finds segment of length 2	Valid
Column 1	`1 1`	Finds segment of length 2	Valid
Column 2	`0 1`	Finds segment of length 1	Valid
Column 3	`1 0`	Finds segment of length 1	Valid
Coloring	`(i+j)%3+1`	Produces colors	Valid

The output coloring is one possible result:

YES
1 2 0 1
2 3 1 0

Every segment is at most length 3, so the cyclic pattern never repeats inside a segment.

Now consider:

1
1 5 3
1 1 1 1 1

Step	Position	Action	Current check result
Row 0	five consecutive tiles	Segment length becomes 4, then 5	Invalid

The algorithm stops after detecting that five tiles require five different colors while only three exist.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(nm)	Each cell is inspected a constant number of times during row scans, column scans, and construction.
Space	O(nm)	The grid is stored so it can be checked and then colored.

The total number of cells over all test cases is at most 100000, so the linear scan easily fits within the limits.

Test Cases

import sys
import io

def solve_input(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    input = sys.stdin.readline

    t = int(input())
    ans = []

    for _ in range(t):
        n, m, k = map(int, input().split())
        wall = [list(map(int, input().split())) for _ in range(n)]

        ok = True

        for i in range(n):
            cur = 0
            for j in range(m):
                if wall[i][j]:
                    cur += 1
                    if cur > k:
                        ok = False
                else:
                    cur = 0

        for j in range(m):
            cur = 0
            for i in range(n):
                if wall[i][j]:
                    cur += 1
                    if cur > k:
                        ok = False
                else:
                    cur = 0

        if not ok:
            ans.append("NO")
        else:
            ans.append("YES")
            for i in range(n):
                row = []
                for j in range(m):
                    if wall[i][j] == 0:
                        row.append("0")
                    else:
                        row.append(str((i + j) % k + 1))
                ans.append(" ".join(row))

    return "\n".join(ans)

assert solve_input("""2
4 3 2
0 1 0
1 1 0
1 0 1
0 1 0
3 4 2
0 1 0 1
1 0 1 1
1 1 1 0
""").splitlines()[0] == "YES"

assert solve_input("""1
1 1 1
1
""") == "YES\n1"

assert solve_input("""1
5 1 3
1
1
1
1
1
""") == "NO"

assert solve_input("""1
2 5 5
1 1 1 1 1
1 1 1 1 1
""").splitlines()[0] == "YES"

assert solve_input("""1
3 3 1
1 0 1
0 1 0
1 0 1
""").splitlines()[0] == "YES"

Test input	Expected output	What it validates
Single tile with one color	`YES`	Minimum size case
Five tiles in one column with three colors	`NO`	Detecting an impossible long segment
Full grid with `k` equal to segment length	`YES`	Maximum allowed segment length
One-color grid separated by lamps	`YES`	Lamps correctly split segments

Edge Cases

For the impossible column example:

The column scan increases cur from 1 to 5. When it reaches 4, the algorithm already knows the segment is longer than the available color count. It outputs NO without attempting a construction.

For a row split by lamps:

1
1 5 3
1 1 0 1 1

The row scan sees a segment of length 2, then a lamp resets the counter, then another segment of length 2. Both are valid because each needs only two different colors. The construction assigns colors independently through the same global formula, and the lamp prevents the two parts from interacting.

For the largest possible segment that still works:

1
1 5 5
1 1 1 1 1

The scan accepts the row because its length equals k. The colors become:

1 2 3 4 5

No value repeats, so the wall is correctly colored.