Welcome to our weekly review of AtCoder Beginner Contest 433! In this series, we explore programming problems from AtCoder’s weekly contests and turn difficult concepts into simple solutions.

A - Happy Birthday! 4

If you are new to AtCoder, check out the official getting-started page!

Problem. You are given positive integers $X$, $Y$ and $Z$. It is guaranteed that $Z\geq2$.

Takahashi is $X$ years old and Aoki is $Y$ years old. Each year, Takahashi and Aoki's ages increase by $1$ simultaneously. Is there any moment in the future (including this year) that Takahashi's age is exactly $Z$ times Aoki's age?

Solution. We can simulate Takahashi and Aoki's ages and check if Takahashi's age ever becomes $Z$ times Aoki's age. If $X\lt YZ$, the quantity on the left-hand side only increases by $1$ while the right-hand side increases by $Z$ every year, so the inequality will remain true in the future. Thus, we can halt the simulation once $X\lt YZ$ becomes true.

Implementation in Python

main.py

X, Y, Z = (int(x) for x in input().split())
while X >= Y * Z:
    if X == Y * Z:
        print("Yes")
        exit()
    X += 1
    Y += 1
print("No")

We can also directly compute the moment when Takahashi's age becomes exactly $Z$ times Aoki's age. Suppose this happens $t$ years from now. Since $X+t=Z(Y+t)$, we can rearrange this equation into $t=\dfrac{X-YZ}{Z-1}$. We can use this to check whether $t$ is a non-negative integer that gives the number of years into the future.

Implementation in Python

main.py

X, Y, Z = (int(x) for x in input().split())
t = (X - Y * Z) // (Z - 1)
if t >= 0 and X + t == Z * (Y + t):
    print("Yes")
else:
    print("No")

B - Nearest Taller

Problem. There are $N$ people standing in a row from left to right. In this problem, person $i$ refers to the $i$-th person from the left. The height of person $i$ is $A_i$.

For each $i\in\{1, 2, \dots, N\}$, determine if there exists a person who is taller than person $i$ and stands to the left of person $i$. If no such person exists, output $-1$. If they exist, output the number of the one standing closest to person $i$.

Solution. We can use a nested for-loop to find the answer for each $i$.

Implementation in Python

main.py

N = int(input())
A = [int(x) for x in input().split()]
result = [-1] * N
for i in range(N):
    for j in range(i - 1, -1, -1):
        if A[j] > A[i]:
            result[i] = j + 1
            break
print(*result)

Alternatively, we can store candidates in a stack. For each person $i$, we are only interested in the closest person whose height is greater than $A_i$, so we can iteratively pop candidates from the stack until a valid candidate is found or the stack becomes empty. Since this avoids using a nested loop, it is more efficient than the previous solution. The time complexity is $\mathcal{O}(N)$.

Implementation in Python

main.py

N = int(input())
A = [int(x) for x in input().split()]

result = [-1] * N
stack = []
for i in range(N):
    while stack and A[stack[-1]] <= A[i]:
        stack.pop()
    if stack:
        result[i] = stack[-1] + 1
    stack.append(i)

print(*result)

C - 1122 Substring 2

Problem. A string $T$ is called a 1122-string if it satisfies the following conditions:

$T$ is a non-empty string consisting of digits.
$|T|$ is even.
The first $|T|/2$ digits are the same.
The last $|T|/2$ digits are the same.
The first digit is $1$ less than the last digit.

You are given a digit string $S$ of length at most $10^6$. Find the number of substrings that are 1122-strings. In particular, two substrings should be counted separately if they are extracted from different positions.

Solution. Since a 1122-string consists of two blocks of identical digits, we can compress adjacent, identical digits into a single block to simplify the problem. This method of compressing strings into blocks (called "runs") of identical characters is known as run-length encoding. Any 1122-substring must consist of exactly two blocks of identical digits, so in the RLE form, it corresponds to using the same number of characters from each of two adjacent runs. We can iterate through every pair of adjacent chunks, giving us a solution in $\mathcal{O}(|S|)$ time.

Implementation in Python

main.py

S = input()
rle = []
for c in S:
    if rle and rle[-1][0] == c:
        rle[-1][1] += 1
    else:
        rle.append([c, 1])

result = 0
for i in range(len(rle) - 1):
    if int(rle[i][0]) + 1 == int(rle[i + 1][0]):
        result += min(rle[i][1], rle[i + 1][1])

print(result)

D - 183183

Problem. Given positive integers $x$ and $y$, let $z$ be the string obtained by concatenating $x$ and $y$ in their decimal notation as strings, and define $f(x, y)$ as the numerical value of $z$ in decimal notation. For example, $f(3, 14)=314$ and $f(100, 3)=1003$.

You are given a sequence $A=(A_1, A_2, \dots, A_N)$ of $N\leq2\times10^5$ positive integers at most $10^9$. You are also given a positive integer $M\leq10^9$. Find the number of pairs of integers $(i, j)$ such that

$1\leq i, j\leq N$, and
$f(A_i, A_j)$ is a multiple of M.

Solution. For positive integers $x$ and $y$, we can rephrase the definition of $f$ as $f(x, y)=x\cdot10^{d(y)}+y$ where $d(y)$ is the number of digits of $y$ in decimal notation. If $f(x, y)$ is a multiple of $M$, then $x\cdot10^{d(y)}+y\equiv0\pmod{M}$, giving us

$$y\equiv -x\cdot10^{d(y)}\pmod{M}.$$

Thus, if we precompute the values of $-A_i\cdot10^d\mod M$ for every possible number of digits $d$, we can iterate over indices $j$ to efficiently count the number of elements $A_i$ for which $f(A_i, A_j)$ is a multiple of $M$.

A solution that stores precomputed values in a hash map runs in $\mathcal{O}(N\log\max{A})$ time.

Implementation in Python

main.py

from collections import Counter

N, M = (int(x) for x in input().split())
A = [int(x) for x in input().split()]

A_pow10 = [-a for a in A]
counters = []
for k in range(10):
    for i in range(N):
        A_pow10[i] *= 10
        A_pow10[i] %= M
    counters.append(Counter(A_pow10))

result = 0
for a in A:
    result += counters[len(str(a)) - 1][a % M]

print(result)

E - Max Matrix 2

Problem. You are given positive integers $N, M$. You are also given a sequence of $N$ positive integers $X=(X_1, \dots, X_N)$ and a sequence of $M$ positive integers $Y=(Y_1, \dots, Y_M)$. Every element in $X$ and $Y$ is at most $N\times M$.

Determine if there exists an $N\times M$ matrix $A$ that satisfies the following conditions, and output such a matrix if it exists:

$1\leq A_{ij}\leq N\times M$;
All elements $A_{ij}$ are distinct;
The maximum element in the $i$-th row of $A$ is $X_i$ for all $i\in\{1, \dots, N\}$; and
The maximum element in the $j$-th column of $A$ is $Y_j$ for all $j\in\{1, \dots, M\}$.

You are given $T$ test cases. The sum of $N\times M$ in all test cases is at most $2\times10^5$.

Solution. Consider the following key observations:

For every element in the matrix, $A_{ij}$ is at most $\min\{X_i, Y_j\}$.
In each row $i$, exactly one element is equal to $X_i$.
In each column $j$, exactly one element is equal to $Y_j$.
If $X_i=Y_j$, the element $A_{ij}$ must be equal to $X_i$. (Otherwise, we would need two copies of that value, violating uniqueness.)

Based on these observations, we can use a greedy algorithm to construct a solution.

Elements $A_{ij}$ where $X_i=Y_j$ can only take this unique value, so we will immediately assign values to them first. Also mark row $i$, column $j$, and the value $X_i$ as "assigned". If the value is already assigned, no solution exists because it violates uniqueness.
We now consider other elements $A_{ij}$ where $X_i\neq Y_j$. Since larger values are more restricted by the upper bounds in observation 1, we will greedily assign larger values to elements $A_{ij}$ in the order of decreasing $\min\{X_i, Y_j\}$. For every $A_{ij}$:
- Let $v=\min\{X_i, Y_j\}$. If $v$ is not yet assigned, set $A_{ij}=v$ and mark row $i$, column $j$ and value $v$ as "assigned".
- If the value $v$ is marked assigned but the row or column is not, a solution does not exist because it violates observations 2 and 3.
- If row $i$, column $j$, and value $v$ are all assigned, we will defer the assignment and push $A_{ij}$ into an unassigned queue.
If there are unassigned elements in the queue, we fill them with unassigned values in descending order. For each element $A_{ij}$ in the queue, let $v$ be the greatest unassigned value and assign $A_{ij}=v$. If $v\gt X_i$ or $v\gt Y_j$, no solution exists because every remaining, unassigned element in the queue has an upper bound at most $\min\{X_i, Y_j\}$ due to assignment order in step 2.

This solution has a time complexity of $\mathcal{O}(NM\log NM)$ (or $\mathcal{O}(NM)$ with bucket sort).

Implementation in Python

main.py

def solve():
    N, M = (int(x) for x in input().split())
    X = [int(x) for x in input().split()]
    Y = [int(x) for x in input().split()]

    assignment_order = []
    for i in range(N):
        for j in range(M):
            assignment_order.append((X[i] == Y[j], min(X[i], Y[j]), i, j))
    assignment_order.sort(reverse=True)

    A = [[0] * M for _ in range(N)]
    assigned_row = [False] * N
    assigned_col = [False] * M
    assigned_val = [False] * (N * M + 1)
    unassigned_queue = []
    for _, val, i, j in assignment_order:
        if assigned_row[i] and assigned_col[j]:
            unassigned_queue.append((i, j))
            continue
        if assigned_val[val]:
            print("No")
            return
        A[i][j] = val
        assigned_row[i] = True
        assigned_col[j] = True
        assigned_val[val] = True

    val = N * M
    for i, j in unassigned_queue:
        while assigned_val[val]:
            val -= 1
        if val > X[i] or val > Y[j]:
            print("No")
            return
        A[i][j] = val
        val -= 1

    print("Yes")
    for row in A:
        print(*row)

for _ in range(int(input())):
    solve()

F - 1122 Subsequence 2

Problem. A string $T$ is called a 1122-string if it satisfies the following conditions. (The definition is the same as in Problem C.)

$T$ is a non-empty string consisting of digits.
$|T|$ is even.
The first $|T|/2$ digits are the same.
The last $|T|/2$ digits are the same.
The first digit is $1$ less than the last digit.

You are given a digit string $S$ of length at most $10^6$. Find the number of subsequences that are 1122-strings, and output the result modulo $998244353$. In particular, two subsequences should be counted separately if they are extracted from different positions.

Solution. Since there are only $10$ digits, we only need to consider nine types of 1122-strings from 00...11 to 88...99. The answer is the sum of possible subsequences for each of the nine types. To simplify the problem, we will assume $S$ only consists of 0 and 1 digits.

Whereas substrings are contiguous parts of a string, a subsequence may not necessarily be contiguous. Nevertheless, the order of digits in a subsequence may not change. Thus, if we choose $S_i=$0 to be a part of the subsequence $T$ for some index $i$, we can also include 0's that appear before $i$ and 1's that appear after $i$. In the following part, we will focus on counting the number of subsequences in which $S_i$ is chosen to be the last 0.

Suppose $S$ contains $\ell$ occurrences of 0 up to index $i$, and $r$ occurrences of 1 after that. Since $S_i$ is chosen to be the last 0, the earlier 0's only have $\ell-1$ choices left. The number of ways to construct a 1122-string of length $2k$ is $\displaystyle\binom{\ell-1}{k-1}\binom{r}{k}$. The sum of this quantity for all possible $k$ gives us the total number of valid subsequences: $\displaystyle\sum_{k=1}^{\min\{\ell, r\}}\binom{\ell-1}{k-1}\binom{r}{k}$.

Thus, the answer is the sum of the above quantity for all indices $i$ for which $S_i=$0. However, if $\dfrac{|S|}2$ digits are 0 and $\dfrac{|S|}2$ digits are 1, the summation requires $\mathcal{O}(|S|^2)$ computational time in the worst case.

A more efficient solution involves a clever transformation of the summation. Indeed, searching the above quantity on WolframAlpha reveals the following equation:

$$\sum_{k=1}^{\min\{\ell, r\}}\binom{\ell-1}{k-1}\binom{r}{k}=\binom{\ell+r-1}{\ell}$$

The general form of this equation is known as Vandermonde's identity. It has a rather simple yet elegant proof involving combinatorics. Using this identity, we can rewrite our earlier solution to run in $\Theta(|S|)$ time with precomputation of binomial coefficients.

Proof

Suppose $\ell\leq r$. We will show that

$$\sum_{k=1}^\ell\binom{\ell-1}{k-1}\binom{r}{k}=\binom{\ell+r-1}{\ell}.$$

By symmetry, $\displaystyle\binom{\ell-1}{k-1}=\binom{\ell-1}{(\ell-1)-(k-1)}=\binom{\ell-1}{\ell-k}$.

Suppose there are $\ell-1$ white balls and $r$ black balls. The quantity $\displaystyle\binom{\ell-1}{\ell-k}\binom{r}{k}$ is the number of ways to choose exactly $\ell-k$ white balls and $k$ black balls, and its summation from $k=1$ to $\ell$ gives the total number of ways to choose $(\ell-k)+k=\ell$ balls. Since this is identical to choosing $\ell$ balls from all $(\ell-1)+r$ balls without distinguishing between their colors, the quantity is also equal to $\displaystyle\binom{\ell+r-1}{\ell}$. The two quantities on each side of the equation are hence identical.

The proof for $\ell\gt r$ is similar.

Implementation in Python

main.py

S = input()
MOD = 998244353

def comb(n, k):
    ... # omitted for brevity

occurrences = [[] for _ in range(10)]
for i in range(len(S)):
    occurrences[int(S[i])].append(i)

result = 0
for d in range(9):
    j = 0
    for i in range(len(occurrences[d])):
        while j < len(occurrences[d + 1]) and occurrences[d][i] > occurrences[d + 1][j]:
            j += 1
        if j == len(occurrences[d + 1]):
            break
        l = i + 1
        r = len(occurrences[d + 1]) - j
        result += comb(l + r - 1, l)
        result %= MOD

print(result)

Well done!

That's another contest down, but don't forget to come back for future digests! See you next time, and watch out for those nasty TLEs! ⏱️🚀