Welcome to our weekly review of AtCoder Beginner Contest 428, a beginner-to-intermediate level programming contest hosted by AtCoder.

We'll highlight the main ideas behind each problem and present sample solutions.

A - Grandma's Footsteps

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

Problem. Takahashi is playing a game at school. Immediately after the bell rings, Takahashi repeats the following actions:

Run at a speed of $S$ meters per second for $A$ seconds, then remain stationary for $B$ seconds.

After $X$ seconds since the bell rings, what will the total distance run by Takahashi be (in meters)?

Solution. Since Takahashi's action is periodic, we can separately calculate the distance run by Takahashi during each cycle and add up the results. Direct calculation is also possible by dividing the total time into complete cycles of $A+B$ seconds.

Implementation in Python (simulation)

main.py

S, A, B, X = (int(x) for x in input().split())
distance = 0
while X > 0:
    distance += S * min(A, X)
    X -= A + B
print(distance)

Implementation in Python (direct calculation)

main.py

S, A, B, X = (int(x) for x in input().split())
print((X // (A + B) * A + min(X % (A + B), A)) * S)

B - Most Frequent Substrings

Problem. You are given a string $S$ consisting of $N$ lowercase English letters.

Let $x$ be the maximum number of times a substring of length $K$ appears in $S$. For example, if $S=$ababa and $K=3$, then $x=2$ because the substring aba appears twice in $S$.

Find $x$ and output all substrings of length $K$ that appear $x$ times in $S$. Output the substrings in lexicographical order.

Solution. We can use a dictionary to count the number of times a substring appears in $S$. Once we determine $x$, we can store valid substrings in an array and sort them in lexicographical order.

Implementation in Python

main.py

from collections import Counter

N, K = (int(x) for x in input().split())
S = input()

num_appearances = Counter()
for i in range(N - K + 1):
    num_appearances[S[i : i + K]] += 1

x = max(num_appearances.values())
print(x)
result = sorted(s for s, c in num_appearances.items() if c == x)
print(*result)

C - Brackets Stack Query

Problem. A string $T$ that consists of only characters ( and ) is called a good bracket sequence if $T$ can be modified into an empty string by repeatedly removing substrings of () from $T$. For example, (), (()()) and the empty string are good bracket sequences, whereas )( and (() are not.

Let $S$ be an empty string. Process $Q\leq8\times10^5$ queries in order:

1 c: Append $c$ to the end of $S$. For this query, $c$ is guaranteed to be either ( or ).
2: Remove the last character of $S$. For this query, $S$ is guaranteed to be non-empty.

After each query, output whether $S$ is a good bracket sequence.

Solution. We can rephrase the definition of a good bracket sequence by replacing ( with the number $1$ and ) with the number $-1$. Given an integer sequence $A=(A_1, \dots, A_N)$ consisting of $1$ and $-1$, define $B$ as its prefix sum sequence such that $B_i=\displaystyle\sum_{j=1}^iA_j$. $A$ is a good bracket sequence if:

$B_i\geq0$ for all $i$, and
$B_N=0$.

Here, the first condition requires that the number of right brackets ) cannot exceed the number of left brackets ( in any prefix of the sequence, and the second condition requires that the entire sequence has an equal amount of left and right brackets. For example, the string (()()) maps to $B=(1, 2, 1, 2, 1, 0)$. The prefix sum sequence is always non-negative, and the final value is zero, so the original string is also a good bracket sequence.

In the original problem, since each query only modifies the end of the string, we can store $B_i$ using a stack and record the number of negative entries to check whether both conditions are met after each query. The implementation below also stores $B_0=0$ for the empty string.

Our solution processes all $Q$ queries in $\Theta(Q)$ time.

Implementation in Python

main.py

prefix_sum = [0]
num_negative = 0

for _ in range(int(input())):
    query = input().split()
    if query[0] == '1':
        c = query[1]
        if c == '(':
            prefix_sum.append(prefix_sum[-1] + 1)
        else:
            prefix_sum.append(prefix_sum[-1] - 1)
        num_negative += prefix_sum[-1] < 0

    else:
        num_negative -= prefix_sum[-1] < 0
        prefix_sum.pop()

    if prefix_sum[-1] == 0 and num_negative == 0:
        print("Yes")
    else:
        print("No")

D - 183184

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 $T\leq3\times10^5$ test cases. In each test case, you are given positive integers $C\leq2\times10^8$ and $D\leq5\times10^9$. Find the number of integers $x$ that satisfy the following conditions:

$1\leq x\leq D$, and
$f(C, C+x)$ is a perfect square.

Solution. Given positive integer $n$, let $d(n)$ be the number of digits of $n$ in decimal notation. We can rephrase the second condition in the problem as:

$f(C, C+x)=C\cdot10^{d(C+x)}+C+x$ is a perfect square.

We can iterate over all possible digit lengths of $C+x$ and count the number of $x$ that satisfy the above equation. When we work with a specific number of digits $d(C+x)$, we can treat $C\cdot10^{d(C+x)}$ like a constant, so we can once again rephrase the problem in simpler terms:

Given positive integers $C$, $C^\prime$ and $d$, what is the number of integers $x$ such that

$C+x$ has $d$ digits, and

$C^\prime+C+x$ is a perfect square?

If $C+x$ has $d$ digits, then $10^{d-1}-1\lt C+x\leq10^d-1$, so

$$C^\prime+10^{d-1}-1\lt C^\prime+C+x\leq C^\prime+10^d-1.$$

We can hence deduce that the number of perfect squares in this interval is

$$\lfloor\sqrt{C^\prime+10^d-1}\rfloor-\lfloor\sqrt{C^\prime+10^{d-1}-1}\rfloor.$$

The original problem has another condition $1\leq x\leq D$, but we can use similar logic to handle edge cases as well.

Since our solution groups cases by the number of digits of $C+x$, the overall time complexity becomes $\mathcal{O}(\log(C+D))$.

Watch out for precision-related issues when computing the square root of a number.

Implementation in Python

main.py

from math import sqrt

def floor_sqrt(x):
    y = int(sqrt(x))
    while y ** 2 > x:
        y -= 1
    while (y + 1) ** 2 <= x:
        y += 1
    return y

for _ in range(int(input())):
    C, D = (int(x) for x in input().split())
    min_digits = len(str(C + 1))
    max_digits = len(str(C + D))
    result = 0

    for d in range(min_digits, max_digits + 1):
        C_PRIME = C * 10 ** d
        if d == min_digits:
            lower_bound = C_PRIME + C + 1
        else:
            lower_bound = C_PRIME + 10 ** (d - 1)
        if d == max_digits:
            upper_bound = C_PRIME + C + D
        else:
            upper_bound = C_PRIME + 10 ** d - 1
        result += floor_sqrt(upper_bound) - floor_sqrt(lower_bound - 1)

    print(result)

E - Farthest Vertex

Problem. You are given an undirected tree with $N\leq5\times10^5$ vertices labeled $1$ to $N$. The $i$-th edge connects vertices $A_i$ and $B_i$.

For each vertex $v=1, \dots, N$, find the vertex with the maximum distance from $v$. If there are multiple vertices with the maximum distance, output the vertex labeled with the largest number.

Solution. If vertex $u$ has the maximum distance from $v$ in the tree, we say that $u$ is one of the farthest vertices from $v$.

The diameter of a tree is the maximum distance between any two vertices in the tree. The following property is useful for solving our problem:

Let $p$ and $q$ be two vertices with distance equal to the diameter. For every vertex $v$, at least one of $p$ and $q$ is one of the farthest vertices from $v$.

We can prove this property by showing that if there exists some vertex $r$ with a greater distance from $v$ than $p$ and $q$ do, then at least one of the paths $p\rightarrow r$ and $q\rightarrow r$ has a distance greater than the diameter, leading to a contradiction.

To find two vertices with a distance equal to the diameter, we can run BFS from an arbitrary vertex $v$ to find a farthest vertex $p$ from $v$, then run BFS from $p$ to find a farthest vertex $q$ from $p$. During both runs, if multiple vertices have the maximum distance, we will break the tie by choosing the vertex labeled with the largest number.

The above solution is sufficient to guarantee that at least one of $p$ and $q$ is the farthest vertex with the greatest label from each vertex. To prove this property, we can follow a similar proof by contradiction to show that if some vertex $r$ has the maximum distance with a greater label, we would have chosen $r$ during one of the BFS runs.

To find $p$ and $q$ and compute their distances to other vertices, we can run BFS three times with a time complexity of $\Theta(N)$.

Implementation in Python

main.py

from collections import deque

N = int(input())
tree = [[] for _ in range(N)]
for _ in range(N - 1):
    A, B = (int(x) - 1 for x in input().split())
    tree[A].append(B)
    tree[B].append(A)

def bfs(source):
    q = deque()
    q.append(source)
    dist = [-1] * N
    dist[source] = 0
    while q:
        u = q.popleft()
        for v in tree[u]:
            if dist[v] != -1:
                continue
            q.append(v)
            dist[v] = dist[u] + 1
    farthest = N - 1
    for v in range(N - 2, -1, -1):
        if dist[v] > dist[farthest]:
            farthest = v
    return dist, farthest

_, p = bfs(0)
dist_p, q = bfs(p)
dist_q, _ = bfs(q)

if p < q:
    p, q = q, p
    dist_p, dist_q = dist_q, dist_p

for v in range(N):
    if dist_p[v] >= dist_q[v]:
        print(p + 1)
    else:
        print(q + 1)

F - Pyramid Alignment

Problem. There are $N\leq2\times10^5$ intervals on a number line labeled from $1$ to $N$. Initially, the left endpoint of interval $i$ is at coordinate $0$ and the right endpoint is at $W_i$. It is guaranteed that $1\leq W_1\lt\dots\lt W_N\leq10^9$.

Process the following $Q\leq2\times10^5$ queries:

1 v: For all intervals $i\lt v$, translate interval $i$ such that its left endpoint matches the left endpoint of interval $v$.
2 v: For all intervals $i\lt v$, translate interval $i$ such that its right endpoint matches the right endpoint of interval $v$.
3 x: Output the number of intervals that contain coordinate $x+\frac12$.

Solution. As the title of the problem suggests, the intervals form a pyramid where the length of each interval is greater than the previous one.

Observe that interval $i$ is nested in interval $i+1$ for all $i\in[1, N)$ after any sequence of operations. More precisely, interval $i$ is always a proper subset of interval $i+1$. Hence, if interval $i$ contains coordinate $x+\frac12$, we can deduce that intervals $i+1, \dots, N$ also contain the coordinate. On the other hand, if interval $i$ does not contain the coordinate, neither do intervals $1, \dots, i-1$. This property is important, as it enables us to use binary search to answer queries on the number of intervals containing a coordinate.

Since the other queries only update consecutive intervals, a natural solution is to use a segment tree with lazy propagation. The AtCoder Library has an implementation <atcoder/lazysegtree> that also supports binary search. Read the documentation here for more details.

With AtCoder Library's implementation, each query can be processed in $\mathcal{O}(\log N)$ time, and the time complexity of our solution is $\mathcal{O}(N+Q\log N)$.

Since updates only affect the prefix of the sequence of intervals, an alternative solution is to use a deque to store blocks of intervals with the same endpoints. While this solution may look inefficient at first, the number of blocks to be inserted and removed is bounded by the number of queries, so the time complexity of updates after all $Q$ queries is $\mathcal{O}(Q)$. The official editorial uses this solution (but with a stack instead of a deque) to achieve $\mathcal{O}(Q\log N)$ runtime.

Implementation in C++

main.cpp

#include <iostream>
#include <vector>
#include <atcoder/lazysegtree>

struct S {
    int anchor = 0, length = 0;
    bool align_left = true, is_id = true;
};

S op(S a, S b) { return b.is_id ? a : b; }
S e() { return {}; }

struct F {
    int anchor = 0;
    bool align_left = true, is_id = true;
};

S mapping(F f, S x) { return f.is_id ? x : S{f.anchor, x.length, f.align_left, false}; }
F composition(F f, F g) { return f.is_id ? g : f; }
F id() { return {}; }

int main() {
    int N;
    std::cin >> N;
    std::vector<S> vec(N);
    for (int i = 0; i < N; i++) {
        std::cin >> vec[i].length;
        vec[i].is_id = false;
    }
    atcoder::lazy_segtree<S, op, e, F, mapping, composition, id> seg(vec);

    int Q;
    std::cin >> Q;
    while (Q--) {
        int t, v, x;
        std::cin >> t;
        if (t == 1) {
            std::cin >> v;
            v--;
            auto [anchor, length, align_left, _] = seg.get(v);
            if (!align_left) anchor -= length;
            seg.apply(0, v, {anchor, true, false});
        }
        else if (t == 2) {
            std::cin >> v;
            v--;
            auto [anchor, length, align_left, _] = seg.get(v);
            if (align_left) anchor += length;
            seg.apply(0, v, {anchor, false, false});
        }
        else {
            std::cin >> x;
            int r = seg.max_right(0, [x](S s) {
                auto [anchor, length, align_left, is_id] = s;
                if (is_id) return true;
                if (align_left) return !(anchor <= x && x < anchor + length);
                return !(anchor - length <= x && x < anchor);
            });
            std::cout << N - r << std::endl;
        }
    }
}

And that's it!

It's been a really long while since the last review, but thank you for your time to be back again. Stay safe out there and watch out for those nasty TLEs! ⏱️🚀