3831 - Medians

Cerința

Se dă un vector cu n elemente. Să se determine numărul de secvențe care au medianul valorilor egal cu k.

Date de intrare

Fișierul de intrare input.txt contine pe prima linie un număr N reprezentând numărul de elemente din vector și un număr k cu semnificația din enunț. Pe a doua linie se află N elemente , elementele vectorului.

Date de ieșire

Fișierul de ieșire output.txt contine pe prima linie răspunsul.

Restricții și precizări

• N ≤ 100.000, K ≤ 1.000.000.000

Exemplul 1

input.txt:

2 2

1 2

output.txt:

1

Exemplul 2

input.txt:

3 5

5 1 5

output.txt:

3

Exemplul 3

input.txt:

3 99999999999

5 1 5

Output:

Invalid input constraints.

Rezolvare

<syntaxhighlight lang="python3" line="1"> MOD = 666013 Nmax = 100001 Add = 100000

def verify_constraints(n, k):

   if not (1 <= n <= 100000 and 1 <= k <= 1000000000):
       print("Invalid input constraints.")
       exit()

def Update(pos, val):

   global aib
   i = pos
   while i < Nmax * 2:
       aib[i] += val
       i += i & -i

def Query(pos):

   global aib
   i = pos
   sum_ = 0
   while i > 0:
       sum_ += aib[i]
       i -= i & -i
   return sum_

def Fx(k):

   global n, V, aib

   aib = [0] * (Nmax * 2)

   S = [0] * (n + 1)

   for i in range(1, n + 1):
       if V[i] > k:
           S[i] = -1
       else:
           S[i] = 1

   for i in range(2, n + 1):
       S[i] += S[i - 1]

   for i in range(1, n + 1):
       Update(S[i] + Add, 1)

   s = 0

   for i in range(1, n + 1):
       s += Query(S[i] + Add - 1)

       if S[i] < 0:
           s += 1

       Update(S[i] + Add, -1)

   return s

with open("input.txt", "r") as fin, open("output.txt", "w") as fout:

   n, k = map(int, fin.readline().split())
   verify_constraints(n,k)
   V = [0] + list(map(int, fin.readline().split()))

   sk = Fx(k)
   skk = Fx(k - 1)

   fout.write(str(abs(sk - skk)))

</syntaxhighlight>