Welcome back! Let's jump right in. This week, we will review AtCoder Beginner Contest 440.

In this series, we explore programming problems from AtCoder’s weekly contests and break down complex concepts into simple solutions!

A - Octave

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

Problem. The frequency of a sound doubles each time its pitch is raised by an octave.

You are given two integers $1\leq X\leq444$ and $1\leq Y\leq 3$. If the frequency of a sound is $X$ hertz, what is its frequency, in hertz, if its pitch is raised by $Y$ octaves?

Solution. Problem A from last week also dealt with powers of $2$, and we discussed different ways of evaluating an exponent. For this problem, we can use bit shifts to concisely express an integer multiplied by a power of $2$.

Implementation in Python

main.py

X, Y = (int(x) for x in input().split())
print(X << Y)

B - Trifecta

Problem. $N$ horses, numbered $1$ to $N$, competed in a race. The $i$-th horse reached the goal $T_i$ seconds after the race began. It is guaranteed that $T_i$ are distinct. Find the numbers of the horses that finished first, second, and third, respectively.

Solution. We can sort the $N$ horses by their finishing time to find the winners.

Implementation in Python

main.py

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

horses_by_time = sorted((T[i], i + 1) for i in range(N))
for j in range(3):
    print(horses_by_time[j][1])

C - Striped Horse

Problem. There are $N\leq2\times10^5$ tiles on a straight line. Initially, all tiles are painted white. The cost of painting the $i$-th tile black is $1\leq C_i\leq10^9$. You are also given a positive integer $W\leq2\times10^5$.

Ringo wants to perform the following operations and paint the tiles in a zebra-like pattern:

Choose a positive integer $x$.
For all integers $1\leq i\leq N$, if the remainder of $(i+x)$ divided by $2W$ is less than $W$, paint the $i$-th tile black.

Find the minimum cost for Ringo to paint the tiles in a zebra-like pattern.

You are given $T$ test cases. The sum of $N$ in all test cases is at most $2\times10^5$, and the sum of $W$ in all test cases is at most $2\times10^5$.

Solution. Observe that Ringo's operations always result in a cycle of $W$ white tiles followed by $W$ black tiles regardless of $x$. Thus, tiles that are $2W$ units away from each other should always have the same color. We can hence categorize tiles by the remainder of their index divided by $2W$ and find the total cost for each group. From here, we no longer need to consider individual tiles.

Since the $W$ black tiles must form a consecutive segment in the $2W$-length cycle, we can use a prefix sum array to efficiently compute the cost of painting any segment black. The time complexity of our solution is $\Theta(N+W)$.

Implementation in Python

main.py

for _ in range(int(input())):
    N, W = (int(x) for x in input().split())
    C = [int(x) for x in input().split()]

    costs = [0] * (W * 2)
    for i in range(N):
        costs[i % (W * 2)] += C[i]

    prefix_sum = [0] * (W * 4 + 1)
    for i in range(W * 2):
        prefix_sum[i + 1] = prefix_sum[i] + costs[i]
    for i in range(W * 2):
        prefix_sum[i + 1 + W * 2] = prefix_sum[i + W * 2] + costs[i]

    print(min(prefix_sum[i + W] - prefix_sum[i] for i in range(W * 2)))

D - Forbidden List 2

Problem. You are given a list $A$ of $N\leq3\times10^5$ distinct integers between $1$ and $10^9$.

You are also given $Q\leq3\times10^5$ queries. For each query $j$, answer the following problem:

You are given two integers $1\leq X_j, Y_j\leq10^9$. Find the $Y_j$-th smallest integer that is at least $X_j$ and does not appear in $A$.

Solution. If the queries had asked for integers that do appear in $A$, we could sort $A$ in ascending order and use binary search to find the solution.

We can do something similar even when the queries ask for integers not present in the list. Consider an infinite list $B$ containing all positive integers that do not appear in $A$. We can divide $B$ into $N+1$ open intervals $(A_i, A_{i+1})$ for $i\in\{0, \dots, N\}$. Here, we can choose $A_0=0$ and $A_{N+1}=2\times10^9+1$ as sentinel values.

Now consider a query in which $X_j$ appears in $A$ as the element $A_i$. We can find the $Y_j$-th smallest integer at least $A_i$ that appears in $B$ by the following operations:

Let $k$ be the index such that the open interval $(A_k, A_{k+1})$ contains the solution to the query. We can find $k$ as the minimum index for which the number of elements in the union of $(A_i, A_{i+1}), \dots, (A_{k-1}, A_k)$ is at most $Y_j$.
Let $s$ be the number of elements in the above union. Since it contains integers between $A_i$ and $A_k$, and excludes the $(k-i+1)$ elements that appear in $A$, we obtain $s=(A_k-A_i+1)-(k-i+1)=(A_k-k)-(A_i-i)$.
The $Y_j$-th smallest integer greater than $A_i$ is therefore the $(Y_j-s)$-th smallest integer greater than $A_k$, so the solution is $A_k+(Y_j-s)$.

We can find $k$ and $s$ in the first two operations via binary search on a prefix sum array. For queries where $X_j\notin A$, we can use binary search to find the greatest $A_i$ less than $X_j$ and adjust $Y_j$, giving us a subproblem that can be solved with the above operations.

The time complexity of our solution is $\mathcal{O}((N+Q)\log N)$.

Implementation in Python

main.py

import bisect

N, Q = (int(x) for x in input().split())
A = [0] + sorted(int(x) for x in input().split()) + [2000000001]
queries = [tuple(int(x) for x in input().split()) for _ in range(Q)]

prefix_sum = [A[i] - i for i in range(N + 2)]
for x, y in queries:
    i = bisect.bisect_right(A, x) - 1
    if x != A[i]:
        y += x - (A[i] + 1)
    k = bisect.bisect_left(prefix_sum, y + prefix_sum[i]) - 1
    s = prefix_sum[k] - prefix_sum[i]
    print(A[k] + y - s)

E - Cookies

Problem. You are given three positive integers: $N\leq50$, $K\leq10^5$, and $X\leq10^5$. There are $N$ types of cookies. The deliciousness of a cookie of type $i$ is $-10^9\leq A_i\leq10^9$. Cookies of the same type are not distinguishable.

There are $\binom{N+K-1}K$ ways to choose $K$ cookies with replacement. For each combination, consider the sum of deliciousness of the $K$ cookies. Let $S_1, S_2, \dots$ be these sums in descending order with duplicates included. Find $S_1, \dots, S_X$. It is guaranteed that there exist at least $X$ different ways of choosing $K$ cookies with replacement.

Solution. We will sort $A_i$ in descending order. By doing so, a combination of $K$ cookies of type $1$ will always have the highest deliciousness. To get the combination with the second-highest deliciousness, we should swap out a cookie of type $1$ for a cookie of type $2$. We can repeat the process of swapping out type-$i$ cookies for type-$j$ cookies $(i\lt j)$ until we find the top $X$ combinations.

Suppose the combination with the $k$-th highest deliciousness consists of $C_i$ cookies of type $i$ for $i\in\{1, \dots, N\}$. Observe that swapping out a cookie of type $i$ for a cookie of type $(i+1)$ results in a deliciousness greater than or equal to swapping for a cookie of types $(i+2), (i+3)$ and so on. Thus, when searching for the $(k+1)$-th highest deliciousness, we only need to consider swapping type-$i$ for type-$(i+1)$ for all $i\in\{1, \dots, N-1\}$.

We can implement a BFS-like solution using a priority queue. Since there are $X$ iterations, and each iteration considers up to $N-1$ additional candidates, the total number of combinations we add to the priority queue is at most $NX$, giving us a solution that runs in $\mathcal{O}(NX\log NX)$ time.

The implementation below uses a set to ensure the priority queue contains only unique combinations. Its time complexity is $\mathcal{O}(N^2X\log NX)$. While this is sufficiently fast in C++, an equivalent Python program would instead result in TLE. The official editorial offers a clever solution that avoids repeated combinations and runs in $\mathcal{O}(N\log N+X\log X)$ time.

Implementation in C++

main.cpp

#include <iostream>
#include <vector>
#include <array>
#include <queue>
#include <set>

using array = std::array<int, 50>;

int main() {
    int N, K, X;
    std::cin >> N >> K >> X;
    std::vector<int64_t> A(N);
    for (int64_t &a : A) std::cin >> a;
    std::sort(A.rbegin(), A.rend());

    std::priority_queue<std::pair<int64_t, array>> pq;
    std::set<array> set;
    array init{};
    init[0] = K;
    pq.emplace(A[0] * K, init);
    set.insert(init);
    while (X--) {
        auto [s, arr] = pq.top();
        pq.pop();
        std::cout << s << std::endl;
        for (int i = 0; i < N - 1; i++) {
            if (arr[i] == 0) continue;
            array narr = arr;
            narr[i]--;
            narr[i + 1]++;
            if (set.contains(narr)) continue;
            set.insert(narr);
            pq.emplace(s - A[i] + A[i + 1], narr);
        }
    }
}

F - Egoism

Problem. There are $N$ horses numbered $1$ to $N$. The $i$-th horse has a mood of $1\leq A_i\leq10^6$ and a tidiness of $B_i\in\{1, 2\}$.

At night, the $N$ horses arrange themselves to bathe one after another. For $j\in\{1, \dots, N\}$, the $j$-th horse taking a bath receives a satisfaction defined as follows:

If $j=1$, the satisfaction is equal to the mood of the horse.
Otherwise, the satisfaction is equal to the product of the mood of the $j$-th horse and the tidiness of the $(j-1)$-th horse.

You are given a query for each of the following $Q\leq2\times10^5$ days. For the $k$-th day, answer the following problem:

Modify the $W_k$-th horse to have a mood of $1\leq X_k\leq10^6$ and a tidiness of $Y_k\in\{1, 2\}$, then compute the maximum total satisfaction that the $N$ horses receive if they decide on an optimal order of bathing that night.

Solution. This problem involves two parts: evaluating the maximum satisfaction and updating the solution for each query. First, we need to find an optimal ordering.

If all horses have a tidiness of $1$, the total satisfaction is the sum of their mood. On the other hand, if all horses have a tidiness of $2$, all horses except the first one can double their satisfaction, so placing a horse with the lowest mood at the front will give the maximum satisfaction.

We now consider a combination of tidiness. Let $m$ be the number of horses with tidiness $2$, and let $S$ be the subset of the top $m$ horses with the highest mood regardless of their tidiness. The upper bound of total satisfaction is $\displaystyle\sum_{i=1}^N A_i+\sum_{j\in S}A_j$.

We now attempt to construct an ordering that achieves this upper bound. When we assign horse $i$ to appear immediately after horse $j$, we will say that horse $i$ is the successor of horse $j$. First, we will temporarily order the horses with tidiness $2$ in ascending order of mood. For each horse with tidiness $2$ that is also in subset $S$, we will assign it to be the successor of the current horse in front of it (*). Next, for the remaining horses in subset $S$ that have tidiness $1$, we will assign each of them to be the successor of a horse with tidiness $2$ that does not have a successor yet. Finally, we will order the remaining horses with tidiness $1$ to appear at the front.

Since each horse in $S$ must be the successor of another horse with tidiness $2$, the assignment marked (*) fails if all horses in $S$ have a tidiness of $2$. In this case, we should compromise by swapping out a horse with the lowest mood in $S$ for a horse with the highest mood not in $S$. The maximum achievable satisfaction becomes

$$\sum_{i=1}^N A_i+\sum_{j\in S}A_j-\min_{j\in S}A_j+\max_{j^\prime\notin S}A_{j^\prime}.$$

Now that we know how to obtain the maximum satisfaction, we should also consider what happens when we update the parameters of a horse in each query. Since the above calculations rely on the maximum and minimum values of subsets, we can use sorted multisets to maintain the mood of horses and keep track of the total mood in each subset. An appropriate implementation should only update a constant number of elements per query, giving us a time complexity of $\Theta((N+Q)\log N)$.

Implementation in C++

main.cpp

#include <iostream>
#include <vector>
#include <tuple>
#include <set>

int main() {
    int N, Q;
    std::cin >> N >> Q;
    std::vector<std::pair<int64_t, int>> horses(N);
    for (auto &[a, b] : horses) std::cin >> a >> b;
    std::vector<std::tuple<int, int64_t, int>> queries(Q);
    for (auto &[w, x, y] : queries) std::cin >> w >> x >> y, w--;

    int64_t sum_mood = 0, sum_high = 0;
    std::multiset<int64_t, std::greater<>> ms_b1, ms_low;
    std::multiset<int64_t> ms_b2, ms_high;
    auto insert_mood = [&](int64_t a, int b) {
        sum_mood += a;
        if (b == 1) ms_b1.insert(a);
        else ms_b2.insert(a);
        sum_high += a;
        ms_high.insert(a);
    };
    auto erase_mood = [&](int64_t a, int b) {
        sum_mood -= a;
        if (b == 1) ms_b1.erase(ms_b1.find(a));
        else ms_b2.erase(ms_b2.find(a));
        if (ms_high.size() && *ms_high.begin() <= a) {
            sum_high -= a;
            ms_high.erase(ms_high.find(a));
        } else {
            ms_low.erase(ms_low.find(a));
        }
    };
    auto update_multisets = [&]() {
        while (ms_high.size() < ms_b2.size()) {
            auto L = ms_low.begin();
            sum_high += *L;
            ms_high.insert(*L);
            ms_low.erase(L);
        }
        while (ms_high.size() > ms_b2.size()) {
            auto H = ms_high.begin();
            sum_high -= *H;
            ms_low.insert(*H);
            ms_high.erase(H);
        }
        while (
            ms_high.size()
            && ms_low.size()
            && *ms_high.begin() < *ms_low.begin()
        ) {
            auto H = ms_high.begin();
            auto L = ms_low.begin();
            sum_high += *L - *H;
            ms_high.insert(*L);
            ms_low.insert(*H);
            ms_high.erase(H);
            ms_low.erase(L);
        }
    };

    for (auto [a, b] : horses) {
        insert_mood(a, b);
    }
    update_multisets();

    for (auto [w, x, y] : queries) {
        auto [a, b] = horses[w];
        erase_mood(a, b);
        horses[w] = {x, y};
        insert_mood(x, y);
        update_multisets();

        int64_t result;
        if (ms_b1.empty()) result = sum_mood * 2 - *ms_b2.begin();
        else if (ms_b2.empty()) result = sum_mood;
        else if (*ms_b2.begin() < *ms_high.begin()) result = sum_mood + sum_high;
        else result = sum_mood + sum_high - *ms_b2.begin() + *ms_b1.begin();
        std::cout << result << std::endl;
    }
}

A related problem on maintaining minimum and maximum values of a multiset

ARC210B - Remove Median Operations

G - Haunted House

Problem. You are about to challenge a haunted house. The haunted house has $F\leq10$ floors, and each floor is represented by a grid with $H\leq500$ rows and $W\leq500$ columns. Let $(k, i, j)$ represent the cell located $k$ floors from the highest level, $i$ rows from the top and $j$ columns from the left.

Each room is represented by a character $S_{k,i,j}$:

If $S_{k,i,j}$ is #, cell $(k, i, j)$ is a wall cell.
Otherwise, cell $(k, i, j)$ is an empty cell, and $S_{k,i,j}$ is a digit between 0 and 9 indicating the number of coins in the cell.

You can move from an empty cell $(k, i, j)$ to another empty cell $(k^\prime, i^\prime, j^\prime)$ as long as:

$k=k^\prime$ and $|i-i^\prime|+|j-j^\prime|=1$;
$k+1=k^\prime$, $i=i^\prime$ and $j=j^\prime$; or
$k-1=k^\prime$, $i=i^\prime$, $j=j^\prime$, and cell $(k, i, j)$ contains a ladder.

Note that you may not move outside the grid.

You are given $Q\leq10^5$ queries. Solve each query $q$ independently:

If you start from cell $(G_q, A_q, B_q)$, what is the maximum number of coins in the cells you visit if you can choose any cell and place a ladder there?

Solution. We can model the grid as a directed graph with $F\cdot H\cdot W$ vertices. Excluding transitions that require a ladder, for each valid move from cell $u$ to $v$ in the original problem, we will add an edge from $u$ to $v$ in our graph. We can partition the vertices into strongly connected components (SCC) and decompose the graph into a directed acyclic graph (DAG). For the rest of this solution, we will ignore the original grid and only refer to vertices in this DAG.

We will define the score of vertex $u$ as the maximum number of coins visited in a path starting from $u$. Since every DAG has a topological ordering, we can use dynamic programming to find the score for all vertices. Consider the following variables:

dp0[u] is the score of $u$ without using a ladder.
dp1[u] is the score of $u$ when a ladder is placed on vertex $u$.
dp2[u] is the score of $u$ when a ladder can be placed anywhere in the DAG.

We can precompute these variables and use dp2 to answer all $Q$ queries in $\Theta(Q)$ time.

Computing dp0 is straightforward: by visiting vertices in a reverse topological ordering, each vertex $u$ can choose from a child v and maximize dp0[u] as the sum of the maximum dp0[v] and the number of coins in $u$.

We can also compute dp1 in reverse topological ordering. For each edge $(u, v)$ in the DAG, we can place a ladder on vertex $v$ and temporarily connect $u$ and $v$ in both directions. In an optimal path, we always start from $v$ and reach $u$ via the ladder, then pick a child $v^\prime$ of $u$ and explore the rest of the DAG from $v^\prime$. If $v=v^\prime$, the maximum number of coins is the sum of dp0[v] and the number of coins in $u$. Otherwise, the maximum number of coins is the sum of dp0[u] and the number of coins in $v$. Since there are two possibilities, we need to revisit our definition so that dp0[u] tracks the top two scores obtained by choosing two different outgoing edges.

Finally, we can compute dp2[v] similarly from dp1. For each edge $(u, v)$ in the DAG, we can either place a ladder on vertex $v$ to get the score of dp1[v], or we can ignore the ladder for now, choose a child $w$ of $v$, and solve for $(v, w)$ recursively. If the solution to $(v, w)$ does not place a ladder on $w$, we can set dp2[v] as the sum of dp2[w] and the number of coins in $v$. If the solution places a ladder on $w$ and chooses a different parent $v^\prime$ such that $v\neq v^\prime$, we get the same sum of dp2[w] and the number of coins in $v$. However, if the solution places a ladder on $w$ and reaches $v$ again, we get a different result that simply sets dp2[v] as dp1[w]. Thus, we should also track the top two scores for dp1 obtained by choosing two different incoming edges.

Since we can calculate each step in constant time, we have a complete solution that runs in $\Theta(FHW+Q)$ time. Make sure to handle cases where a vertex does not have outgoing or incoming edges.

The following implementation uses the AtCoder Library for SCC decomposition in $\Theta(FHW)$ time. Instead of recursing, the solution updates dp1 in place when computing dp2. It separately uses a sorting algorithm to remove duplicate edges in the DAG, so its time complexity is $\mathcal{O}(FHW\log HW+Q)$.

Implementation in C++

main.cpp

#include <iostream>
#include <vector>
#include <algorithm>
#include <tuple>
#include <atcoder/scc>

int main() {
    int F, H, W;
    std::cin >> F >> H >> W;
    std::vector S(F, std::vector<std::string>(H));
    for (int k = 0; k < F; k++) for (int i = 0; i < H; i++) std::cin >> S[k][i];

    atcoder::scc_graph scc_graph(F * H * W);
    auto encode = [&](int k, int i, int j) { return k * H * W + i * W + j; };
    auto decode = [&](int x) { return std::tuple{x / H / W, x / W % H, x % W}; };
    for (int k = 0; k < F; k++) {
        for (int i = 0; i < H; i++) {
            for (int j = 0; j < W; j++) {
                if (S[k][i][j] == '#') continue;
                if (k < F - 1 && S[k + 1][i][j] != '#') {
                    scc_graph.add_edge(encode(k, i, j), encode(k + 1, i, j));
                }
                if (i < H - 1 && S[k][i + 1][j] != '#') {
                    scc_graph.add_edge(encode(k, i, j), encode(k, i + 1, j));
                    scc_graph.add_edge(encode(k, i + 1, j), encode(k, i, j));
                }
                if (j < W - 1 && S[k][i][j + 1] != '#') {
                    scc_graph.add_edge(encode(k, i, j), encode(k, i, j + 1));
                    scc_graph.add_edge(encode(k, i, j + 1), encode(k, i, j));
                }
            }
        }
    }

    std::vector<std::vector<int>> scc = scc_graph.scc();
    int N = scc.size();
    std::vector<int> num_coins(N);
    std::vector scc_id(F, std::vector(H, std::vector<int>(W)));
    for (int u = 0; u < N; u++) {
        for (int x : scc[u]) {
            auto [k, i, j] = decode(x);
            scc_id[k][i][j] = u;
            num_coins[u] += S[k][i][j] - '0';
        }
    }

    std::vector<std::vector<int>> dag(N), rev(N);
    for (int k = 0; k < F - 1; k++) {
        for (int i = 0; i < H; i++) {
            for (int j = 0; j < W; j++) {
                if (S[k][i][j] == '#' || S[k + 1][i][j] == '#') continue;
                dag[scc_id[k][i][j]].push_back(scc_id[k + 1][i][j]);
                rev[scc_id[k + 1][i][j]].push_back(scc_id[k][i][j]);
            }
        }
    }
    for (int u = 0; u < N; u++) {
        std::sort(dag[u].begin(), dag[u].end());
        dag[u].erase(std::unique(dag[u].begin(), dag[u].end()), dag[u].end());
        std::sort(rev[u].begin(), rev[u].end());
        rev[u].erase(std::unique(rev[u].begin(), rev[u].end()), rev[u].end());
    }

    std::vector<std::vector<std::pair<int, int>>> dp0(N);
    for (int u = N - 1; u >= 0; u--) {
        dp0[u].emplace_back(num_coins[u], u);
        for (int v : dag[u]) {
            dp0[u].emplace_back(dp0[v][0].first + num_coins[u], v);
        }
        std::sort(dp0[u].rbegin(), dp0[u].rend());
        if (dp0[u].size() > 2u) dp0[u].resize(2);
    }

    std::vector<std::vector<std::pair<int, int>>> dp1(N);
    for (int u = N - 1; u >= 0; u--) {
        dp1[u].emplace_back(dp0[u][0].first, -1);
        for (int v : dag[u]) {
            for (auto [score, dest] : dp0[u]) {
                if (v == dest) dp1[v].emplace_back(score, u);
                else dp1[v].emplace_back(score + num_coins[v], u);
            }
        }
    }

    std::vector<int> dp2(N);
    for (int v = N - 1; v >= 0; v--) {
        std::sort(dp1[v].rbegin(), dp1[v].rend());
        if (dp1[v].size() > 2) {
            for (unsigned i = 1; i < dp1[v].size(); i++) {
                if (dp1[v][i].second == dp1[v][0].second) continue;
                std::swap(dp1[v][i], dp1[v][1]);
                break;
            }
            dp1[v].resize(2);
        }
        dp2[v] = dp1[v][0].first;
        for (int u : rev[v]) {
            for (auto [score, source] : dp1[v]) {
                if (u == source) dp1[u].emplace_back(score, -1);
                else dp1[u].emplace_back(score + num_coins[u], -1);
            }
        }
    }

    int Q;
    std::cin >> Q;
    while (Q--) {
        int g, a, b;
        std::cin >> g >> a >> b;
        g--, a--, b--;
        std::cout << dp2[scc_id[g][a][b]] << std::endl;
    }
}

Look! We're done!

We went through all seven problems in the contest this week. You've earned yourself a pat on the back! See you next time, and watch out for those nasty TLEs! ⏱️🚀