已经改的跟题解几乎一样了,样例能过然后MLE?
查看原帖
已经改的跟题解几乎一样了,样例能过然后MLE?
224558
JackMerryYoung楼主2023/7/7 15:58

Rt, 调了一个下午了,救救孩子吧。

#include <bits/stdc++.h>
// #define int long long
using namespace std;

int N, M;
int w[100005], neww[100005];

struct Edge {
    int nxt, to;
} edge[200005];

int cnt, head[100005];

void add(int u, int v)
{
    ++ cnt;
    edge[cnt].nxt = head[u];
    edge[cnt].to = v;
    head[u] = cnt;
}

int father[100005], depth[100005], heavyson[100005], sz[100005];

void dfs1(int u, int fa)
{
    father[u] = fa;
    depth[u] = depth[fa] + 1;
    sz[u] = 1;
    for(int i = head[u]; i; i = edge[i].nxt)
    {
        int v = edge[i].to;
        if(v == fa) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if(sz[heavyson[u]] < sz[v]) heavyson[u] = v;
    }
}

int dfncnt, dfn[100005], linktop[100005];

void dfs2(int u, int lktp)
{
    dfn[u] = ++ dfncnt;
    neww[dfn[u]] = w[u];
    linktop[u] = lktp;
    if(!heavyson[u]) return;
    dfs2(heavyson[u], lktp);
    for(int i = head[u]; i; i = edge[i].nxt)
    {
        int v = edge[i].to;
        if(v == father[u] || v == heavyson[u]) continue;
        dfs2(v, v);
    }
}

#define lson (p << 1)
#define rson ((p << 1) | 1)

struct SegmentTree {
    int l, r, left_end_color, right_end_color, segmentcnt, let_lazytag;
    public:
        SegmentTree()
        {
            this->l = 0ll, this->r = 0ll, this->left_end_color = 0ll, this->right_end_color = 0ll, this->segmentcnt = 0ll, let_lazytag = 0ll;
        }
} tree[100000 << 2 + 5];

SegmentTree lquery, rquery;

SegmentTree merge_query(SegmentTree a, SegmentTree b)
{
    if(!a.segmentcnt) return b;
    if(!b.segmentcnt) return a;
    SegmentTree res;
    res.left_end_color = a.left_end_color, res.right_end_color = b.right_end_color;
    res.segmentcnt = a.segmentcnt + b.segmentcnt;
    if(a.right_end_color == b.left_end_color) -- res.segmentcnt;
    return res;
}

void merge(int p)
{
    tree[p].left_end_color = tree[lson].left_end_color, tree[p].right_end_color = tree[rson].right_end_color;
    tree[p].segmentcnt = tree[lson].segmentcnt + tree[rson].segmentcnt;
    if(tree[lson].right_end_color == tree[rson].left_end_color) -- tree[p].segmentcnt;
}

void build(int l, int r, int p)
{
    tree[p].l = l, tree[p].r = r;
    if(l == r)
    {
        tree[p].segmentcnt = 1;
        tree[p].left_end_color = tree[p].right_end_color = neww[l];
        return;
    }

    int mid = (l + r) >> 1;
    build(l, mid, lson);
    build(mid + 1, r, rson);
    merge(p);
}

void pushdown(int p)
{
    tree[lson].segmentcnt = tree[rson].segmentcnt = 1;
    tree[lson].left_end_color = tree[lson].right_end_color = tree[rson].left_end_color = tree[rson].right_end_color = tree[p].let_lazytag;
    tree[p].let_lazytag = 0;
}

void modify_let(int lll, int rrr, int p, int val)
{
    if(lll <= tree[p].l && rrr >= tree[p].r)
    {
        tree[p].segmentcnt = 1;
        tree[p].let_lazytag = tree[p].left_end_color = tree[p].right_end_color = val;
        return;
    }

    if(tree[p].let_lazytag) pushdown(p);
    int mid = (tree[p].l + tree[p].r) >> 1;
    if(lll <= mid) modify_let(lll, rrr, lson, val);
    if(rrr > mid) modify_let(lll, rrr, rson, val);
    merge(p);
}

void modify_let_simplepath(int x, int y, int val)
{
    while(linktop[x] != linktop[y])
    {
        if(depth[linktop[x]] < depth[linktop[y]]) swap(x, y);
        modify_let(dfn[linktop[x]], dfn[x], 1, val);
        x = father[linktop[x]];
    }
    if(depth[x] > depth[y]) swap(x, y);
    modify_let(dfn[x], dfn[y], 1, val);
}

SegmentTree query_segmentcnt(int lll, int rrr, int p) 
{
	if(tree[p].l >= lll && tree[p].r <= rrr) return tree[p];
	if(tree[p].let_lazytag) pushdown(p);
	int mid = (tree[p].l + tree[p].r) >> 1;
	if(rrr <= mid) return query_segmentcnt(lll, rrr, lson);
	if(lll > mid) return query_segmentcnt(lll, rrr, rson);
	return merge_query(query_segmentcnt(lll, rrr, lson), query_segmentcnt(lll, rrr, rson));	
}

int query_segmentcnt_simplepath(int x, int y)
{
    lquery = SegmentTree(), rquery = SegmentTree();
    while(linktop[x] != linktop[y])
    {
        if(depth[linktop[x]] < depth[linktop[y]]) 
        {
            rquery = merge_query(query_segmentcnt(dfn[linktop[y]], dfn[y], 1), rquery);
            y = father[linktop[y]];
        }
        else
        {
            lquery = merge_query(query_segmentcnt(dfn[linktop[x]], dfn[x], 1), lquery);
            x = father[linktop[x]];
        }
    }
    if(depth[x] < depth[y]) rquery = merge_query(query_segmentcnt(dfn[x], dfn[y], 1), rquery);
    else lquery = merge_query(query_segmentcnt(dfn[y], dfn[x], 1), lquery);
    swap(lquery.left_end_color, lquery.right_end_color);
    return merge_query(lquery, rquery).segmentcnt;
}

signed main()
{
    scanf("%d %d", &N, &M);
    for(int i = 1; i <= N; ++ i)
        scanf("%d", &w[i]);
    
    for(int i = 1; i < N; ++ i)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        add(u, v);
        add(v, u);
    }

    dfs1(1, 0);
    dfs2(1, 1);
    build(1, N, 1);
    for(int i = 1; i <= M; ++ i)
    {
        char op;
        cin >> op;
        if(op == 'C')
        {
            int a, b, c;
            scanf("%d %d %d", &a, &b, &c);
            modify_let_simplepath(a, b, c);
        }
        if(op == 'Q')
        {
            int a, b;
            scanf("%d %d", &a, &b);
            printf("%d\n", query_segmentcnt_simplepath(a, b));
        }
    }
    return 0;
}
2023/7/7 15:58
加载中...