USACO DEC08 Problem 'sec' Analysis

by Spencer Liang

For a string s and an integer k, let |s| denote the length of s and let s(k) denote the prefix of s of length k. We say two strings s and t are a match if one of the following two conditions holds:

• |s| < |t| and t(|s|) = s
  
• |t| <= |s| and s(|t|) = t
  

We're given two sets A, B of binary strings and for each string b in B, we're asked to find the number of strings a in A such that a, b match. To do so efficiently, we first build a modified trie T on A. Note that each node v in the trie represents the prefix p(v) of some string in A (the path from the root to v spells out p(v)). In addition to pointers to its children, we also have each node v store

• end(v) = the number of strings a in A such that a = p(v)
• part(v) = the number of strings a in A such that a(|p(v)|) = p(v)

Now assume we are given a string b in B. Let v be the last vertex visited when searching for b in T. We consider the two cases for string matching separately. In the first case, for each prefix b(k) from k = 1 to |b|-1, we want the number of strings in A such that a = b(k). In other words, this is just the sum of end(v') for all vertices v' on the path from the root to v, not including v if b is in T. In the second case, we want the number of strings in A such that a(|b|) = b, or simply part(v) if b is in T and 0 otherwise. So querying takes linear time as well.

#include <stdio.h>
#include <iostream>
using namespace std;
#define REP(i,n) for(int i=0;i<n;i++)

struct node {
    int part, end;
    node *l, *r;
} start;

int main() {
    FILE *fin = fopen("sec.in", "r"), *fout = fopen("sec.out", "w");
    int M, N; fscanf(fin, "%d %d", &M, &N);
    REP(i, M) {
        int b; fscanf(fin, "%d", &b);
        node *at = &start;
        REP(i, b) {
            int bit; fscanf(fin, "%d", &bit);
            if (bit == 0) {
                if (at->l == NULL) at->l = new node();
                at = at->l;
            }
            else {
                if (at->r == NULL) at->r = new node();
                at = at->r;
            }
            at->part++;
        }
        at->end++;
    }

    REP(i, N) {
        int b; fscanf(fin, "%d", &b);
        int matches = 0;
        node *at = &start;
        REP(i, b) {
            int bit; fscanf(fin, "%d", &bit);
            if (at == NULL) continue;
            matches += at->end;
            if (bit == 0) at = at->l;
            else at = at->r;
        }
        if (at != NULL) matches += at->part;
        fprintf(fout, "%d\n", matches);
    }
    return 0;
}