USACO 2020 US Open Contest, Gold Problem 1. Haircut

Problem Link

Time Complexity: O(NlogN)

Algorithm used: Segment Trees

In this problem, the lengths of N hairs are given. For each j = 0, j = 1, ..., j = N-1, reduce all hair lengths greater than j to j. Then find the number of inversions. An inversions is a pair of hairs such that the first hair has a greater length than the second, and precedes the second hair in the input.

We could directly simulate this process by reducing the hairs for every j, and then count the inversions in N² or NlogN time. But, this method would be too slow to pass completely. 

Instead, we need a quick way to know the number of pairs of a and b, such that a > b and a precedes b in the input aka the number of inversions. We do this by maintaining two arrays: one for storing the frequency of the hairs, and another for storing the the number of inversions in terms of b.

The reason why we store the number of inversions in terms of b instead of a is because when a hair is reduced, some of the inversions for that hair become invalid, since that hairs length becomes smaller than its pairs length -> a will be equal to j (length hair is reduced to), but b may still be = j. Example, (5, 2) is a pair, and the hair with length 5 is reduced to 2. The new pair becomes (2, 2), which is invalid, since a must be more than b. Some pairs may remain valid. The validity of pairs for each a is difficult to find out quickly without storing and then looping through all pairs and checking.

The number of inversions for b is the number of hairs that have a length greater than be and that precede b in the input. Use a segment tree or a binary indexed tree to store this number for each hair length. Do this by scanning the hairs in the order they appear in the input and update the tree. After updating the tree for the ith hair, update the pairs by finding the sum of the frequencies greater than b, which gives the number of pairs such that a > b, until the ith hair. 

When hairs are reduced to length j, then it is guaranteed that all pairs stored in terms of b, will be valid if b < j, since if b < j and b < a, then a must be < j. So, for each j, the number of inversions is the sum of pairs from b = 0, b = 1, ..., b < j.

Code

#include <bits/stdc++.h>

using namespace std;

int n, seg_size;

vector <long long> freq_seg_tree, inversions, compressed_nums;

int get_id(int x){

    return lower_bound(begin(compressed_nums), end(compressed_nums), x) - begin(compressed_nums);

}

long long sum(int a, int b){

    a += seg_size;

    b += seg_size;

    long long sum = 0;

    while (a <= b){

        if (a % 2 == 1)

            sum += freq_seg_tree[a++];

        if (b % 2 == 0)

            sum += freq_seg_tree[b--];

        a /= 2;

        b /= 2;

    }

    return sum;

}

void update(int a, long long x){

    a += seg_size;

    freq_seg_tree[a] = x;

    for (a /= 2; a >= 1; a /= 2)

        freq_seg_tree[a] = freq_seg_tree[a * 2] + freq_seg_tree[a * 2 + 1];

}

int main()

{

    freopen("haircut.in", "r", stdin);

    freopen("haircut.out", "w", stdout);

    cin >> n;

    vector <int> nums(n);

    compressed_nums.resize(n);

    for (int i = 0; i < n; i++){

        cin >> nums[i];

        compressed_nums[i] = nums[i];

    }

    // compress sorted nums

    sort(begin(compressed_nums), end(compressed_nums));

    compressed_nums.erase(unique(begin(compressed_nums), end(compressed_nums)),

     end(compressed_nums));

   

    /* build segment tree

       for each hair i, store the number of hairs that precede i,

       and whose length is greater than length of hair i

    */

    seg_size = pow(2, ceil(log2(compressed_nums.size())));

    freq_seg_tree.resize(2 * seg_size);

    inversions.resize(compressed_nums.size());

    for (int i = 0; i < n; i++){

        update(get_id(nums[i]), freq_seg_tree[get_id(nums[i]) + seg_size] + 1);

        long long s = sum(get_id(nums[i]+1), compressed_nums.size()-1);

        inversions[get_id(nums[i])] += s;

    }

    long long ans = 0, prev_pos = -1;

    for (int j = 0; j < n; j++){

        if (j == 0){

            cout << 0 << '\n';

            continue;

        }

        int x = lower_bound(begin(compressed_nums), end(compressed_nums), j) -

         begin(compressed_nums);

        x--;

        if (prev_pos != x){

            prev_pos = x;

            ans += inversions[x];

        }

        cout << ans << '\n';

    }

    return 0;

}