Codeforces Round #291 Div. 2 D - R2D2 and Droid Army

Problem Statement:
514D - R2D2 and Droid Army


Solution:
The problem statement is a bit hard to understand, but other than that the problem is excellent. I love the idea of combining binary search with segment tree / Fenwick tree!

The strategy is as follows: We want to know, for each i-th droid, what is the longest segment [j, i] possible such that we can kill all j-th to i-th droid using at most k shots (of course j \(\leq \) i). To find j, we can use binary search on [0, i] droids. If we need more than k shots to kill all droids in [j,i], we must increase j, since any choices of j less than the current j will also need more than k shots. Otherwise, if we need at most k shots to kill droids in [j,i], we can decrease j, since any choices of j bigger than the current j will also need at most k shots. What an awesome property.



So what remains is to compute the minimum number of shots needed to kill all droids in [j,i]. For this purpose we can use segment tree or Fenwick tree. The tree will have m fields, one field for each type of "detail" (whatever it means), and each field keeps track of the largest number of details of that type possessed by droids in [j,i]. After building such a tree, computing the number of shots needed for each [j,i] will take \(O(\lg{N})\). Hence in total we will take \(O(N\lg{N}^2)\). (Btw there is an O(N) solution by simply iterating through the elements! The approach is explained in the official editorial).

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;

int segtree[400005][6];
int cache[400005][6];
int a[100005][6];
int N,M,K;

void build(int p, int L, int R) {
    if(L==R){
        for(int i=0;i<M;++i){
            segtree[p][i] = a[L][i];
        }
        return;
    }
    int mid = (L+R)/2;
    build(2*p,L,mid);
    build(2*p+1,mid+1,R);
    for(int i=0;i<M;++i){
        segtree[p][i]=max(segtree[2*p][i],segtree[2*p+1][i]);
    }
}

int rmq(int p,int L,int R,int S,int T){
    if(R<S || T<L)return -1;
    if(S<=L && R<=T){
        for(int i=0;i<M;++i){
            cache[p][i]=segtree[p][i];
        }
        return p;
    }
    int mid=(L+R)/2;
    int left = rmq(2*p,L,mid,S,T);
    int right = rmq(2*p+1,mid+1,R,S,T);
    if(left==-1)return right;
    if(right==-1)return left;
    for(int i=0;i<M;++i){
        cache[p][i] = max(cache[left][i],cache[right][i]);
    }
    return p;
}

int main(){
    scanf("%d%d%d",&N,&M,&K);
    for(int i=0;i<N;++i){
        for(int j=0;j<M;++j){
            scanf("%d",&a[i][j]);
        }
    }
    build(1,0,N-1);
    int ans = -1;
    int b[6] = {0,0,0,0,0,0};
    for(int i=0;i<N;++i){
        int lo=0,hi=i,mid;
        while(lo<=hi){
            mid=(lo+hi)/2;
            int val = 0;
            int x = rmq(1,0,N-1,mid,i);
            for(int j=0;j<M;++j){
                val += cache[x][j];
            }
            if(val <= K){
                hi=mid-1;
                if(ans < i-mid+1){
                    ans = i-mid+1;
                    for(int j=0;j<M;++j)b[j]=cache[x][j];
                }
            }else{
                lo=mid+1;
            }
        }
    }
    for(int j=0;j<M;++j){
        if(j!=0)printf(" ");
        printf("%d",b[j]);
    }
    printf("\n");
    return 0;
}