Problem p03703 - Generation 2

Orig Description

Score : 600 points
Problem StatementYou are given an integer sequence of length N, a = {a_1, a_2, …, a_N}, and an integer K.
a has N(N+1)/2 non-empty contiguous subsequences, {a_l, a_{l+1}, …, a_r} (1 ≤ l ≤ r ≤ N). Among them, how many have an arithmetic mean that is greater than or equal to K?
Constraints
All input values are integers.
1 ≤ N ≤ 2 \times 10^5
1 ≤ K ≤ 10^9
1 ≤ a_i ≤ 10^9
InputInput is given from Standard Input in the following format:
N K
a_1
a_2
:
a_N
OutputPrint the number of the non-empty contiguous subsequences with an arithmetic mean that is greater than or equal to K.
Sample Input 13 6
7
5
7
Sample Output 15
All the non-empty contiguous subsequences of a are listed below:
{a_1} = {7}
{a_1, a_2} = {7, 5}
{a_1, a_2, a_3} = {7, 5, 7}
{a_2} = {5}
{a_2, a_3} = {5, 7}
{a_3} = {7}
Their means are 7, 6, 19/3, 5, 6 and 7, respectively, and five among them are 6 or greater. Note that {a_1} and {a_3} are indistinguishable by the values of their elements, but we count them individually.
Sample Input 21 2
1
Sample Output 20
Sample Input 37 26
10
20
30
40
30
20
10
Sample Output 313

Extracted Specification

An integer \( N \) (1 ≤ \( N \) ≤ 2 × 10^5), representing some quantity or size.
An integer \( K \) (1 ≤ \( K \) ≤ 10^9), representing some threshold or parameter.
A list of \( N \) integers \( A \), where each integer \( A[i] \) (1 ≤ \( A[i] \) ≤ 10^9).

### Example Input:

```
3 6
7
5
7
```

### Function Signature:
Write a function f(N, K, A) that takes in the input.
```python
def f(N: int, K: int, A: List[int]):
    '''
    N: an integer
    K: an integer
    A: a list of integers
    '''
```

Test Case Generator

Certainly! Here is a test case generator for the given problem statement:

```python
import random

def tcgen():
    # Generate N within given constraints
    N = random.randint(1, 2 * 10**5)
    
    # Generate K within given constraints
    K = random.randint(1, 10**9)
    
    # Generate sequence a of length N with each element in the range [1, 10^9]
    a = [random.randint(1, 10**9) for _ in range(N)]
    
    return N, K, a

# Example usage:
N, K, a = tcgen()
print(N, K)
print('\n'.join(map(str, a)))
```

This script generates random test cases according to the problem constraints:

1. `N` is an integer between 1 and \(2 \times 10^5\).
2. `K` is an integer between 1 and \(10^9\).
3. The sequence `a` contains `N` integers, each between 1 and \(10^9\).

The example usage at the end demonstrates how to generate and print a test case. Each test case is generated randomly, ensuring a wide range of possible scenarios to test the solution.

Extract Arguments

def extract_arguments(fh):
    N, K = map(int, fh.readline().strip().split())
    A = [int(fh.readline().strip()) for _ in range(N)]
    return N, A, K