1789 - 3 Secv

Cerința

Se consideră un șir de n elemente numere naturale. O subsecvență se definește ca o succesiune de elemente ale șirului luate de pe poziții consecutive.

De exemplu şirul 1 4 3 5 are 10 subsecvenţe:

• 4 subsecvențe de lungime 1
  • 1
  • 4
  • 3
  • 5
• 3 subsecvențe de lungime 2
  • 1 4
  • 4 3
  • 3 5
• 2 subsecvențe de lungime 3
  • 1 4 3
  • 4 3 5
• 1 subsecvenţă de lungime 4
  • 1 4 3 5

Să se scrie un program care pentru un şir cunoscut determină pentru câte subsecvenţe ale şirului suma elementelor care le alcătuiesc are un anumit rest dat la împărţirea cu 3.

Date de intrare

Fișierul de intrare 3secv.in conține pe prima linie două numere naturale n şi r, separate prin spaţiu, n reprezentând numărul de elemente din şirul dat, iar r reprezentând restul pentru care trebuie rezolvată cerinţa. Pe linia a doua vor fi n numere naturale separate prin câte un spaţiu, reprezentând elementele şirului dat.

Date de ieșire

Fișierul de ieșire 3secv.out va conține pe prima linie un număr natural reprezentând numărul de subsecvențe ale șirului pentru care suma elementelor dă restul r la împărțirea cu 3.

Restricții și precizări

• 1 ≤ n ≤ 1.000.000
• r ∈ {0,1,2}
• Elementele șirului sunt numere naturale mai mici sau egale cu 32767

Exemplul 1

3secv.in

4 0

1 4 3 5

3secv.out

2

Exemplul 2

3secv.in

4 1

1 4 3 5

3secv.out

4

Exemplul 3

3secv.in

4 2

1 4 3 5

3secv.out

4

Exemplul 4

3secv.in

1000001 1

1 2 3 4 5 6 1000001

3secv.out

Date de intrare invalide.

Rezolvare

#1789 3secv
def is_valid_input(n, r, arr):
  # Verificăm dacă n și r sunt în limitele impuse
  if not (1 <= n <= 1000000 and r in {0, 1, 2}):
    return False
  # Verificăm dacă toate elementele din șir sunt numere naturale mai mici sau egale cu 32767
  if any(not (1 <= x <= 32767) for x in arr):
    return False
  return True

def count_subsequences(n, arr, r):
  count = 0
  for i in range(n):
    for j in range(i, n):
      subsequence_sum = sum(arr[i:j+1])
      if subsequence_sum % 3 == r:
        count += 1
  return count

def main():
  # Citirea datelor de intrare
  with open("3secv.in", "r") as fin:
    n, r = map(int, fin.readline().split())
    arr = list(map(int, fin.readline().split()))

  # Verificarea datelor de intrare
  if not is_valid_input(n, r, arr):
    # Scrierea mesajului de date invalide în fișierul de ieșire
    with open("3secv.out", "w") as fout:
      fout.write("Date de intrare invalide.\n")
    return

  # Calculul numărului de subsecvențe cu suma dată prin restul r la împărțirea cu 3
  result = count_subsequences(n, arr, r)
  
  # Scrierea rezultatului în fișierul de ieșire
  with open("3secv.out", "w") as fout:
    fout.write(str(result) + "\n")

if __name__ == "__main__":
  main()