NAIPC2018 Prefix Free Code(字典树）

问题 E: Prefix Free Code

题目描述

Consider n initial strings of lower case letters, where no initial string is a prefix of any other initial string. Now, consider choosing k of the strings (no string more than once), and concatenating them together. You can make this many such composite strings:
n × (n − 1) × (n − 2) × . . . × (n − k + 1)
Consider sorting all of the composite strings you can get via this process in alphabetical order. You are given a test composite string, which is guaranteed to belong on this list. find the position of this test composite string in the alphabetized list of all composite strings, modulo 10⁹ +7. The first composite string in the list is at position 1.

输入

Each input will consist of a single test case. Note that your program may be run multiple times on different inputs. Each test case will begin with a line with two integers, first n and then k(1 ≤ k ≤ n), where n is the number of initial strings, and k is the number of initial strings you choose to form composite strings. The upper bounds of n and k are limited by the constraints on the strings, in the following paragraphs.
Each of the next n lines will contain a string, which will consist of one or more lower case letters a..z. These are the n initial strings. It is guaranteed that none of the initial strings will be a prefix of any other of the initial strings.
finally, the last line will contain another string, consisting of only lower case letters a..z. This is the test composite string, the position of which in the sorted list you must find. This test composite string is guaranteed to be a concatenation of k unique initial strings.
The sum of the lengths of all input strings, including the test string, will not exceed 10⁶ letters.

输出

Output a single integer, which is the position in the list of sorted composite strings where the test composite string occurs. Output this number modulo 10⁹ + 7.

样例输入

5 3
a
b
c
d
e
cad

样例输出

26

#include <iostream>

#include <stdio.h>
#include <string>
#include <map>
#include <algorithm>
#include <vector>

using namespace std;

const int maxn = 1e6 + 10;
const int mod = 1e9 + 7;
const int inf = 0x3f3f3f3f;
#define ull unsigned long long
#define ll long long
int n, k;
string list[maxn];
map<string, int> Map;
long long ans[maxn];
int ID;

class node {
public:
    bool isWord;
    int rank;
    node **next;
    int myid;

    // Set up an empty node.
    node() {
        isWord = false;
        next = new node *[26];
        for (int i = 0; i < 26; i++) {
            next[i] = nullptr;
        }
        myid = ID++;
    }

    // Insert s (at char k) which has rank myrank.
    void insert(string s, int k, int myrank) {
        // Got to end, set up node.
        if (k == s.length()) {
            isWord = true;
            rank = myrank;
            return;
        }
        // We don't have a node down this path, create it.
        if (next[s[k] - 'a'] == nullptr) {
            next[s[k] - 'a'] = new node();
        }
        // Recursively insert.
        next[s[k] - 'a']->insert(s, k + 1, myrank);
    }
};

node *mytrie;
vector<int> initPerm;

void getPerm(string item) {
    node *ptr = mytrie;
    // Go through the whole string.

    for (int i = 0; i < item.length(); i++) {

        // Update ptr.
        ptr = ptr->next[item[i] - 'a'];

        // If it's a word, add its rank to our permutation and move our ptr back.
        if (ptr->isWord) {
            initPerm.push_back(ptr->rank);
            ptr = mytrie;
        }
    }
}

// This is Travis's bit. Mine works to, but in contest when debugging, we thought that might be
// the issue (it wasn't) and subbed it with this one.
class bit {
public:
    int n;
    int v[maxn];

    int sum(int x) {
        return x < 0 ? 0 : (v[x] + sum((x & (x + 1)) - 1));
    }

    void inc(int p, int x) {
        for (int i = p; i < n; i |= (i + 1)) v[i] += x;
    }

    bit(int N) {
        n = N;
        //v = new int[n];
    }
};

string item;

int main() {
    //freopen("C:\Users\albert\Desktop\7622\prefix-2028.in", "r", stdin);
    //freopen("input.txt", "r", stdin);
    // Read in n, k and list of words, then sort.
    cin >> n >> k;
    if (n == 1 && k == 1) {
        cout << 1 << endl;
        return 0;
    }
    for (int i = 0; i < n; i++) {
        cin >> list[i];
    }
    sort(list, list + n);

    // Store alpha rank of each word.
    for (int i = 0; i < n; i++) {
        Map[list[i]] = i;
    }

    // Store relevant products for rankings.
    ans[0] = 1;
    for (int i = 1; i < k; i++) {
        ans[i] = (ans[i - 1] * (n - k + i)) % mod;
    }

    // Store words in a trie so we can read in the long string quickly.
    mytrie = new node();
    for (int i = 0; i < n; i++) {
        mytrie->insert(list[i], 0, i);
    }

    // Get the long string and the corresponding permutation.
    cin >> item;
    long long res = 0;
    getPerm(item);
    bit *mybit = new bit(n + 1);

    // Calculate rank going through the word.
    for (int i = 0, j = k - 1; i < initPerm.size(); i++, j--) {
        // Get this rank in the permutation.
        int nextItem = initPerm[i];

        // These are how many items before this have already occurred in the perm.
        int sub = mybit->sum(nextItem + 1);

        // This is our ACTUAL ranking of this item of the items left to place.
        int real = nextItem - sub;

        // We add to the ranking the contribution of this element, by multiplying how many items
        // that haven't been placed that come before this one times the appropriate # of permutations.
        res = (res + real * ans[j] + mod) % mod;

        // Update our bit to indicate that we've used this item (using its original rank).
        mybit->inc(nextItem + 1, 1);

    }

    printf("%lld
", (res + 1) % mod);
    return 0;
}