Codeforces Round 1037 Problem C "I will Definitely Make It"

Link: https://codeforces.com/contest/2126/problem/C

Time limit per test: 1 second

Memory limit per test: 256 megabytes

Problem description

You are given n towers, numbered from 1 to n. Tower i has a height of h_i. At time 0, you are on the tower with index k, and the current water level is 1.

Every second, the water level rises by 1 unit. At any moment, if the water level becomes strictly greater than the height of the tower you are on, you perish.

You have a magical ability: at moment x, you can start teleporting from tower i to tower j, which will take abs(h_i − h_j) seconds (abs means the absolute value function). That is, until moment x + abs(h_i − h_j), you will be on tower i, and at moment x + abs(h_i − h_j), you will move to tower j. You can start a new teleportation at the same moment you just arrived at tower j.

For example, if n = k = 4, h = [4,4,4,2], then if you start teleporting from tower 4 to tower 1 at moment 0, the movement will look as follows:

Note that if the height of tower 1 were 5, you would not be able to teleport to it immediately, as you would be submerged at moment 2.

Your goal is to reach any tower with the maximum height before the water covers you.

Determine if this is possible.

Input

Each test consists of several test cases. The first line contains a single integer t (1 ≤ t ≤ 10⁴) — the number of test cases. The description of the test cases follows.

The first line of each test case contains two integers n and k (1 ≤ k ≤ n ≤ 10⁵) — the number of towers and the index of the tower you are initially on.

The second line contains n integers h₁,h₂,…,h_n (1 ≤ h_i ≤ 10⁹) — the heights of the towers.

It is guaranteed that the sum of all n across all test cases does not exceed 10⁵.

Output

For each test case, output one line: “YES”, if you can reach the tower with the maximum height before the water covers you, or “NO” otherwise.

You may output each letter in any case (lowercase or uppercase). For example, the strings “yEs”, “yes”, “Yes”, and “YES” will be accepted as a positive answer.

Example:

Input:

5
5 3
3 2 1 4 5
3 1
1 3 4
4 4
4 4 4 2
6 2
2 3 6 9 1 2
4 2
1 2 5 6

Output:

YES
NO
YES
YES
NO

Note

In the first test case, the only possible path is: 3 → 2 → 1 → 4 → 5.

In the second test case, regardless of the order, it will not be possible to reach the tallest tower.

In the third test case, one of the possible paths is: 4 → 1.

Approach

When thinking about the problem, it might seem the solution could be to just immediately move to the tallest tower. However, because the water is rising each second, you probably will perish before you can teleport.

The solution is to keep teleporting to the shortest tower that is taller than the tower you are currently standing on. Let’s take the example of towers of height 2, 4 and 5. If you start on tower of height 2, immediately trying to teleport to tower of height 5 will cause you to perish (it would take 3 seconds to teleport, which is too long).

Instead, we can first move to tower or height 4, then tower of height 5. This allows us to eventually move to the tallest tower without actualy perishing. This method also allows us to move in the same amount of time as if you would immediately try to teleport to the tallest tower.

This is the best approach because you always want to stay as high as possible, and this method takes the same amount of time as just teleporting to the tallest tower.

Code

#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    int T;
    cin >> T;
    for(int i = 0; i < T; i++) {
        int N, K;
        cin >> N >> K;
        vector<int> H(N);
        int starting_height = 0;
        for(int j = 0; j < N; j++) {
            cin >> H[j];
            if(K == j+1) starting_height = H[j];
        }
        sort(H.begin(), H.end());
        if(starting_height == H[N-1]) {
            cout << "YES\n";
            continue;
        }
        vector<int> difference_of_heights;
        for(int j = 0; j < N; j++) {
            if(j == N-1 || H[j] != H[j+1] && H[j] > starting_height) {
                difference_of_heights.push_back(H[j]);
            }
        }
        int current_water_height = 1;
        int curr_height = starting_height;
        bool works = true;
        for(int j = 0; j < difference_of_heights.size(); j++) {
            current_water_height += (difference_of_heights[j]-curr_height-1);
            if(current_water_height > curr_height) {
                cout << "NO\n";
                works = false;
                break;
            } else {
                current_water_height++;
                curr_height = difference_of_heights[j];
            }
        }
        if(works) {
            cout << "YES\n";
        }
    }
    return 0;
}

Lines 1-5: imported all libraries needed and used the namespace std

Lines 8-9: initializes T (test cases) and inputs them in

Line 10: starts a for loop for the number of test cases

Line 11-12: initializes and inputs in N and K

Lines 13-19: intializes a vector (array) called H, inputs in the heights of all towers, finds the height of the tower we are currently on and sorts the vector H in non-decreasing order

Lines 20-23: if the tower we started with is the tallest tower (or has the same height as the tallest tower), then we can print out 0 and continue with the next test cases

Lines 24-29: finds the differences between heights of consecutive (by height, not by location) towers

Lines 30-46: tries the approach we talked before, makes sure that we do not perish before going onto the next tower, prints NO if at some point we do perish before moving on and prints YES if we do succeed at getting to the tallest tower

Then, the code continues this process from lines 11-46 for all other test cases.

Accepted!