TLE on Train Schedule

# #include <bits/stdc++.h>

# using namespace std;

# int main(){

#     int n;
#     cin >> n;

#     vector<int> a(n), ans(n, 0);
#     for(int &i : a) cin >> i;

#     vector<bool> visited(n, 0);
#     for(int i=0; i<n; i++){
#         if(visited[i]) continue;
#         if(a[i]*a[i] < n){
#             vector<int> count(a[i], 0);
#             for(int j=i, k=0; j<n; j++, k=(k+1)%a[i]){
#                 count[k] += a[i] == a[j];
#                 visited[j] = visited[j] | (a[i] == a[j]); 
#                 ans[j] += count[k];
#             }
#         } else {
#             for(int j=i; j<n; j+=a[i]) ans[j]++;
#         }
#     }

#     for(int i : ans) cout << i << ' ';

#     return 0;
# }

def test():
    n = 100000
    a = "1 " * n
    return str(n) + \
            str(a)[:-1]



if __name__ == "__main__":
    res = test()
    print(res)

Test 1

Verdict: WRONG ANSWER