Train schedule

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


void task5() {
    int n;
    cin >> n;
    vector<int> nums(n);
    for(int i = 0; i < n; i++) cin >> nums[i];

    vector<int> trainsNeeded(n, 0);
    vector<int> longInt;
    vector<int> startDay;
    int sq = sqrt(n) + 1;

    for(int i = 0; i < n; i++) {
        if(nums[i] >= sq) {
            longInt.push_back(nums[i]);
            startDay.push_back(i);
        }
    }

    vector<vector<int>> count(sq);
    for(int i = 0; i < n; i++) {
        if(nums[i] < sq) {
            count[nums[i]].push_back(i);
        }
    }

    for(int interval = 1; interval < sq; interval++) {
        vector<int> add(n + 1, 0);
        for(int start : count[interval]) add[start]++;
        for(int day = 0; day < n; day++) {
            if(day >= interval) add[day] += add[day - interval];
            trainsNeeded[day] += add[day];
        }
    }

    for(size_t i = 0; i < longInt.size(); i++) {
        for(int j = startDay[i]; j < n; j += longInt[i]) trainsNeeded[j]++;
    }

    for(int i = 0; i < n; i++) cout << trainsNeeded[i] << " ";
    cout << endl;
}





void task4(){
    int n; 
    cin >> n;

    for(int i = 0; i < n; i++){
        for(int j = 0; j < n; j++){
            for(int k = 0; k < n; k++){
                if((k + j + i) % 3 == 0) 
                   cout << "R" << " ";
                else if((k+j +i) % 3 == 1)
                    cout << "G" << " ";
                else
                    cout << "B" << " ";
            }
            cout << endl;
        }
    }
}
int main() {
    

task5();
    return 0;

}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: ACCEPTED

Test 7

Verdict: ACCEPTED

Test 8

Verdict: ACCEPTED

Test 9

Verdict: ACCEPTED

Test 10

Verdict: ACCEPTED

Test 11

Verdict: ACCEPTED

Test 12

Verdict: ACCEPTED

Test 13

Verdict: ACCEPTED

Test 14

Verdict: ACCEPTED

Test 15

Verdict: ACCEPTED

Test 16

Verdict: ACCEPTED

Test 17

Verdict: ACCEPTED

Test 18

Verdict: ACCEPTED

Test 19

Verdict: ACCEPTED

Test 20

Verdict: ACCEPTED

Test 21

Verdict: ACCEPTED

Test 22

Verdict: ACCEPTED

Test 23

Verdict: ACCEPTED

Test 24

Verdict: ACCEPTED

Test 25

Verdict: ACCEPTED

Test 26

Verdict: ACCEPTED

Test 27

Verdict: ACCEPTED

Test 28

Verdict: ACCEPTED

Test 29

Verdict: ACCEPTED

Test 30

Verdict: ACCEPTED

Test 31

Verdict: ACCEPTED

Test 32

Verdict: ACCEPTED

Test 33

Verdict: ACCEPTED

Test 34

Verdict: ACCEPTED

Test 35

Verdict: ACCEPTED

Test 36

Verdict: ACCEPTED

Test 37

Verdict: ACCEPTED

Test 38

Verdict: ACCEPTED

Test 39

Verdict: ACCEPTED

Test 40

Verdict: ACCEPTED

Test 41

Verdict: ACCEPTED

Test 42

Verdict: ACCEPTED

Test 43

Verdict: ACCEPTED

Test 44

Verdict: ACCEPTED

Test 45

Verdict: ACCEPTED

Test 46

Verdict: ACCEPTED

Test 47

Verdict: ACCEPTED

Test 48

Verdict: ACCEPTED

Test 49

Verdict: ACCEPTED

Test 50

Verdict: ACCEPTED

Test 51

Verdict: ACCEPTED

Test 52

Verdict: ACCEPTED

Test 53

Verdict: ACCEPTED

Test 54

Verdict: ACCEPTED

Test 55

Verdict: ACCEPTED

Test 56

Verdict: ACCEPTED

Test 57

Verdict: ACCEPTED

Test 58

Verdict: ACCEPTED

Test 59

Verdict: ACCEPTED

Test 60

Verdict: ACCEPTED

Test 61

Verdict: ACCEPTED

Test 62

Verdict: ACCEPTED

Test 63

Verdict: ACCEPTED

Test 64

Verdict: ACCEPTED

Test 65

Verdict: ACCEPTED

Test 66

Verdict: ACCEPTED

Test 67

Verdict: ACCEPTED

Test 68

Verdict: ACCEPTED

Test 69

Verdict: ACCEPTED

Test 70

Verdict: ACCEPTED

Test 71

Verdict: ACCEPTED

Test 72

Verdict: ACCEPTED