Submission #1251176

#	Time	Username	Problem	Language	Result	Execution time	Memory
1251176		MojoLake	Triple Peaks (IOI25_triples)	C++20	51 / 100	2113 ms	905756 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;

vector<vector<int>> easy(vector<int> h) {
    int n = sz(h);

    vector<vector<int>> ret;

    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];
            if (k < n && h[k] == k - i) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        auto case2 = [&]() { // to the left
            int k = i + h[j];
            if (k < n && k != j && h[k] == abs(k - j)) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        auto case3 = [&]() {
            int k = j - h[j];
            if (k >= 0 && k != i && h[k] == k - i) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        case1();
        case2();
        case3();
    }

    return ret;
}

vector<vector<int>> easy_rev(vector<int> h) {
    int n = sz(h);
    vector<vector<int>> ret;

    for (int k = n - 1; k >= 0; --k) {
        int j = k - h[k];

        if (j < 0) continue;

        auto case1 = [&]() {
            int i = j - h[j];
            if (i >= 0 && h[i] == k - j) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        auto case2 = [&]() {
            int i = k - h[j];
            if (i >= 0 && i != j && h[i] == abs(i - j)) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        auto case3 = [&]() {
            int i = j + h[j];
            if (i < n && i != k && h[i] == i - k) {
                vector<int> a = {i, j, k};
                sort(all(a));
                ret.push_back(a);
            }
        };

        case1();
        case2();
        case3();
    }

    return ret;
}

vector<vector<int>> small(int x, vector<int> group, vector<int>& h) {
    int n = sz(h);
    int m = sz(group);
    vector<vector<int>> ret;
    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];
                if (k < n && h[k] == j - i && h[i] == k - j && h[j] == k - i) {
                    vector<int> a = {i, j, k};
                    sort(all(a));
                    ret.push_back(a);
                }
            };

            auto case2 = [&]() {
                int k = j - h[i];
                if (k >= 0 &&  h[k] == j - i && h[i] == j - k && h[j] == abs(k - i)) {
                    vector<int> a = {i, j, k};
                    sort(all(a));
                    ret.push_back(a);
                }
            };

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

            case1();
            case2();
        }
    }
    
    return ret;
}

vector<vector<int>> big(int x, vector<int> group, vector<int>& h) {
    int n = sz(h);
    int m = sz(group);
    vector<vector<int>> ret;
    sort(all(group));

    auto check = [&](vector<int>& v) {
        vector<int> a = {v[1] - v[0], v[2] - v[0], v[2] - v[1]};
        vector<int> b = {h[v[0]], h[v[1]], h[v[2]]};
        sort(all(a));
        sort(all(b));
        return a == b;
    };

    for (int k = 0; k < n; ++k) {
        int z = k - h[k] - x;
        if (z < 0 || z % 2) continue;
        int i = z / 2;
        int j = h[k] + i;

        vector<int> v = {i, j, k};
        sort(all(v));
        if (check(v)) {
            ret.push_back(v);
        }
    }
    return ret;
}

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

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

    int sqr = ceil(sqrt(n));

    vector<vector<int>> ret;

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

        if (sz(group) >= sqr) {
            auto v = big(x, group, h);
            for (auto a : v) ret.push_back(a);
        } else {
            auto v = small(x, group, h);
            for (auto a : v) ret.push_back(a);
        }
    }

    return ret;
}



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

    auto check = [&](vector<int>& v) {
        vector<int> a = {v[1] - v[0], v[2] - v[0], v[2] - v[1]};
        vector<int> b = {h[v[0]], h[v[1]], h[v[2]]};
        sort(all(a));
        sort(all(b));
        return a == b;
    };

    auto a = easy(h);
    auto b = easy_rev(h);
    set<vector<int>> s;
    for (auto v : a) {
        if (check(v)) s.insert(v);
    }
    for (auto v : b) {
        if (check(v)) s.insert(v);
    }

    // for (auto v : s) {
    //     for (int x : v) {
    //         cerr << x << " ";
    //     }
    //     cerr << "\n";
    // }

    ll cnt = sz(s);

    // cerr << "cnt before hard: " << cnt << "\n";

    auto c = hard(h);
    for (auto v : c) {
        if (check(v)) s.insert(v);
    }

    cnt = sz(s);

    return cnt;
}

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

• Subtask 0: 0 / 0
• Subtask 1: 8 / 8
• Subtask 2: 6 / 6
• Subtask 3: 10 / 10
• Subtask 4: 11 / 11
• Subtask 5: 16 / 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