Submission #1251146

#	Time	Username	Problem	Language	Result	Execution time	Memory
1251146		MojoLake	Triple Peaks (IOI25_triples)	C++20	0 / 100	2094 ms	17976 KiB

triples

#include "triples.h"

#include <bits/stdc++.h>


#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()

using namespace std;
using ll = long long;

long long easy(vector<int> h) {
    ll cnt = 0;
    int n = sz(h);

    for (int i = 0; i < n; ++i) {
        int j = i + h[i];

        if (j >= n) continue;
        
        auto case1 = [&]() { // to the right
            int k = j + h[j];
            return (int)(k < n && h[k] == k - i);
        };

        auto case2 = [&]() { // to the left
            int k = i + h[j];
            return (int)(k < n && k != j && h[k] == abs(k - j));
        };

        // cerr << "i: " << i << "\n";
        // cerr << "case1: " << case1() << "\n";
        // cerr << "cas2: " << case2() << "\n";
        cnt += case1();
        cnt += case2();
    }

    return cnt;
}

ll small(int x, vector<int> group, vector<int>& h) {
    int n = sz(h);
    int m = sz(group);
    ll cnt = 0;
    sort(all(group));

    // cerr << "x: " << x << "\n";

    for (int u = 0; u < m - 1; ++u) {
        int i = group[u];
        for (int z = u + 1; z < m; ++z) {
            int j = group[z];
            auto case1 = [&]() {
                int k = j + h[i];
                return (int)(k < n && h[k] == j - i && h[i] == k - j && h[j] == k - i);
            };

            auto case2 = [&]() {
                int k = j - h[i];
                return (int)(k >= 0 &&  h[k] == j - i && h[i] == j - k && h[j] == abs(k - i));
            };

            // cerr << "i and j: " << i << " " << j << "\n";

            cnt += case1() + case2();
        }
    }
    
    return cnt;
}

long long hard(vector<int> h) {
    int n = sz(h);
    map<int, vector<int>> groups;
    ll cnt = 0;

    for (int i = 0; i < n; ++i) {
        groups[h[i] - i].push_back(i);
    }

    int sqr = ceil(sqrt(n));

    for (auto [x, group] : groups) {
        // cerr << "x: " << x << "\n";
        // cerr << "group: ";
        // for (int e : group) {
        //     cerr << e << " ";
        // } 
        // cerr << "\n";

        if (sz(group) >= sqr) {
            // cnt += big(x, group);
            cnt += small(x, group, h);
        } else {
            cnt += small(x, group, h);
        }
    }

    return cnt;
}


long long count_triples(std::vector<int> h) {
    int n = sz(h);

    ll cnt = easy(h);
    reverse(all(h));
    cnt += easy(h);
    reverse(all(h));

    cnt += hard(h);

    return cnt;
}

std::vector<int> construct_range(int M, int K) {
  return {1, 1, 1};
}

• Subtask 0: 0 / 0
• Subtask 1: 0 / 8
• Subtask 2: 0 / 6
• Subtask 3: 0 / 10
• Subtask 4: 0 / 11
• Subtask 5: 0 / 16
• Subtask 6: 0 / 19
• Subtask 7: 0 / 5
• Subtask 8: 0 / 5
• Subtask 9: 0 / 5
• Subtask 10: 0 / 5
• Subtask 11: 0 / 5
• Subtask 12: 0 / 5