grandmother-backpack.cpp

#include <iostream>
#include <set>
#include <bitset>
#include <cstring>
#include <vector>
#include <algorithm>
#include <utility>
using namespace std;

typedef set<int> si;

const int MAX_N = 1000;
int graph[MAX_N][MAX_N];
int N;

void getLoop(int start, vector<int> &visitChain, bitset<MAX_N> &visited, set<int> &loop) {
    visited[start] = true;
    visitChain.push_back(start);

    // follow the graph
    for (int i = 0; i < N; i++) {
        if (graph[start][i]) {
            // make sure it wasn't the previous element
            if (visited[i]) {
                // AHA, a loop!

                vector<int>::iterator it = find(visitChain.begin(), visitChain.end(), i);
                for (; it != visitChain.end(); it++)
                    loop.insert(*it);

                return;
            } else {
                getLoop(i, visitChain, visited, loop);
            }
        }
    }
}

// start is inside the loop for the component when called from outside
int countComponent(int start, bitset<MAX_N> &visited) {
    int count = 1;
    visited[start] = true;
    for (int i = 0; i < N; i++) {
        // reverse edges
        if (graph[i][start] && !visited[i]) {
            count += countComponent(i, visited);
        }
    }
    return count;
}

int main() {
    int K;
    cin >> N >> K;

    for (int i = 0; i <MAX_N;i++)
        memset(graph, 0, MAX_N);

    for (int i = 0; i < N; i++) {
        int pref;
        cin >> pref;
        pref--;
        graph[i][pref] = 1;
    }

    bitset<MAX_N> visited;
    vector<pair<int, int> > compSize;
    // list of loops for each component

    for (int i = 0; i < N; i++) {
        if (visited[i]) continue;
        si loop;
        vector<int> chain;
        bitset<MAX_N> v2;
        getLoop(i, chain, v2, loop);
        int idx = *(loop.begin());
        int s = countComponent(idx, visited);
        compSize.push_back(make_pair(loop.size(), s));
    }

    int numComponents = 0;
    vector<si> compNums (compSize.size());

    for (vector<pair<int, int> >::iterator it = compSize.begin(); it != compSize.end(); it++) {
        compNums[numComponents].insert(0);

        for (int i = (*it).first; i <= (*it).second; i++) {
            if (i <= K)
                compNums[numComponents].insert(i);
        }
        numComponents++;
    }

    vector<set<int> > dp (numComponents);

    for (si::iterator it = compNums[0].begin(); it != compNums[0].end(); it++)
        dp[0].insert(*it);

    // now fill in the rest of the table
    for (int c = 1; c < numComponents; c++) {
        for (si::iterator it = compNums[c].begin(); it != compNums[c].end(); it++) {
            for (si::iterator theyIt = dp[c - 1].begin(); theyIt != dp[c - 1].end(); theyIt++) {
                if (*it + *theyIt <= K)
                    dp[c].insert(*theyIt + *it);
            }
        }
    }

    cout << *(dp[numComponents - 1].rbegin()) << endl;
    return 0;
}