Cod sursă (job #785929)

#include <bits/stdc++.h>

#define f first
#define s second
#define M 1000000
#define ll long long

using namespace std;

const int nMax = 1e5 + 5;

ifstream fin("burlane.in");
ofstream fout("burlane.out");

struct Burlan
{
    int next;
    int lastInBlock;
    int flows;
};

int n, q, nBlocks;
Burlan p[nMax];

void input();
void precalc();
bool isInBlock(int lvl, int next);
void op1();
void op2();
void updateFlows(int x, int y);

int main()
{
    input();
    precalc();
    while(q--) {
        int op;
        fin >> op;
        if(op == 1)
            op1();
        else
            op2();
    }

    return 0;
}

bool isInBlock(int lvl, int next)
{
    return lvl / nBlocks == next / nBlocks;
}

void input()
{
    fin >> n >> q;
    for(int i = 1; i <= n; i++)
        fin >> p[i].next;
}

void precalc()
{
    nBlocks = sqrt(n);
    for(int i = 1; i <= n; i++)
        updateFlows(i, p[i].next);
}

void updateFlows(int x, int y)
{
    if(isInBlock(x, y)) {
        p[x].flows = p[y].flows + 1;
        p[x].lastInBlock = p[y].lastInBlock;
    }
    else {
        p[x].flows = 1;
        p[x].lastInBlock = y;
    }
}

void op1()
{
    int x;
    fin >> x;
    int total = 0;
    while(x) {
        total += p[x].flows;
        x = p[x].lastInBlock;
    }
    fout << total << '\n';
}

void op2()
{
    int x, y;
    fin >> x >> y;
    p[x].next = y;
    updateFlows(x, y);
}