sort-uniq-sort.py

#!/usr/bin/python3

# The alternative script of "cat files|sort|uniq -c|sort -r -n".

import marisa
import sys
import io

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

def parse_args():
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('inputs', nargs='+', help='input file(s)')
    p.add_argument('--limit', type=int, default=1)
    return p.parse_args()

if __name__ == '__main__':
    args = parse_args()
    keyset = marisa.Keyset()
    for file in args.inputs:
        with open(file, encoding='utf-8') as fp:
            for line in fp:
                keyset.push_back(line.rstrip())
    trie = marisa.Trie()
    trie.build(keyset)

    freq = [0] * trie.num_keys()
    agent = marisa.Agent()
    for file in args.inputs:
        with open(file, encoding='utf-8') as fp:
            for line in fp:
                agent.set_query(line.rstrip())
                assert trie.lookup(agent), "%s is not found" % line.rstrip()
                freq[agent.key_id()] += 1

    for i,f in sorted(enumerate(freq), key=lambda pair: -pair[1]):
        if f >= args.limit:
            print("%d\t%s" % (f, trie.reverse_lookup(i)))